From 985c5e91f1267f7d39e6611d36df8b2585fde82d Mon Sep 17 00:00:00 2001
From: sarahwada-stripe <139389030+sarahwada-stripe@users.noreply.github.com>
Date: Tue, 9 Jul 2024 14:11:55 -0700
Subject: [PATCH] DynamoDB leadership election: return leadership success
 status (#687)

* DynamoDB leadership election: return leadership success status

* dynamo db - fix delete tablename
---
 .../mantisrx/extensions/dynamodb/DynamoDBLeaderElector.java | 6 +++++-
 .../java/io/mantisrx/extensions/dynamodb/DynamoDBStore.java | 2 +-
 2 files changed, 6 insertions(+), 2 deletions(-)

diff --git a/mantis-control-plane/mantis-control-plane-dynamodb/src/main/java/io/mantisrx/extensions/dynamodb/DynamoDBLeaderElector.java b/mantis-control-plane/mantis-control-plane-dynamodb/src/main/java/io/mantisrx/extensions/dynamodb/DynamoDBLeaderElector.java
index 0d1315e5c..d6a216b08 100644
--- a/mantis-control-plane/mantis-control-plane-dynamodb/src/main/java/io/mantisrx/extensions/dynamodb/DynamoDBLeaderElector.java
+++ b/mantis-control-plane/mantis-control-plane-dynamodb/src/main/java/io/mantisrx/extensions/dynamodb/DynamoDBLeaderElector.java
@@ -118,8 +118,9 @@ public boolean isLeaderElectorRunning() {
     /**
      * This function will attempt to become the leader at the heartbeat interval of the lockClient. If
      * it becomes the leader it will update the leader data and the thread will stop running.
+     * Returns true if leader election succeeded, false otherwise.
      */
-    private void tryToBecomeLeader() {
+    protected boolean tryToBecomeLeader() {
         final MasterDescription me = leadershipManager.getDescription();
         try {
             log.info("requesting leadership from {}", me.getHostname());
@@ -135,9 +136,12 @@ private void tryToBecomeLeader() {
                 leaderLock = optionalLock.get();
                 shouldLeaderElectorBeRunning.set(false);
                 leadershipManager.becomeLeader();
+                return true;
             }
+            return false;
         } catch (RuntimeException | InterruptedException | JsonProcessingException e) {
             log.error("leader elector task has failed it will restart, if this error is frequent there is likely a problem with DynamoDB based leader election", e);
+            return false;
         } finally {
             isLeaderElectorRunning.set(false);
             if (shouldLeaderElectorBeRunning.get()) {
diff --git a/mantis-control-plane/mantis-control-plane-dynamodb/src/main/java/io/mantisrx/extensions/dynamodb/DynamoDBStore.java b/mantis-control-plane/mantis-control-plane-dynamodb/src/main/java/io/mantisrx/extensions/dynamodb/DynamoDBStore.java
index f6e7da489..629671a49 100644
--- a/mantis-control-plane/mantis-control-plane-dynamodb/src/main/java/io/mantisrx/extensions/dynamodb/DynamoDBStore.java
+++ b/mantis-control-plane/mantis-control-plane-dynamodb/src/main/java/io/mantisrx/extensions/dynamodb/DynamoDBStore.java
@@ -165,7 +165,7 @@ public boolean delete(String tableName, String partitionKey, String secondaryKey
         expressionAttributeValues.put(PK, AttributeValue.builder().s(tableName).build());
         expressionAttributeValues.put(SK, AttributeValue.builder().s(String.format("%s#%s", partitionKey, secondaryKey)).build());
         final DeleteItemRequest request = DeleteItemRequest.builder()
-                .tableName(tableName)
+                .tableName(this.mantisTable)
                 .key(expressionAttributeValues)
                 .build();
         final DeleteItemResponse response = this.client.deleteItem(request);