Register leader election listeners when election is opened.

atomix · Jan 2, 2016 · 5663e46 · 5663e46
1 parent ac220c1
commit 5663e46
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 45 deletions.
diff --git a/coordination/src/main/java/io/atomix/coordination/DistributedLeaderElection.java b/coordination/src/main/java/io/atomix/coordination/DistributedLeaderElection.java
@@ -95,6 +95,7 @@ public class DistributedLeaderElection extends Resource<DistributedLeaderElectio
   public static final ResourceType<DistributedLeaderElection> TYPE = new ResourceType<>(DistributedLeaderElection.class);
 
   private final Set<Consumer<Long>> listeners = Collections.newSetFromMap(new ConcurrentHashMap<>());
+  private volatile long epoch;
 
   public DistributedLeaderElection(CopycatClient client) {
     super(client);
@@ -107,14 +108,15 @@ public ResourceType<DistributedLeaderElection> type() {
 
   @Override
   public CompletableFuture<DistributedLeaderElection> open() {
-    return super.open().thenApply(result -> {
+    return super.open().thenRun(() -> {
       client.<Long>onEvent("elect", epoch -> {
+        this.epoch = epoch;
         for (Consumer<Long> listener : listeners) {
           listener.accept(epoch);
         }
       });
-      return result;
+    }).thenCompose(v -> submit(new LeaderElectionCommands.Listen()))
-    });
+      .thenApply(v -> this);
   }
 
   /**
@@ -145,18 +147,26 @@ public CompletableFuture<DistributedLeaderElection> open() {
    *   }
    * </pre>
    *
-   * @param listener The listener to register.
+   * @param callback The callback to register.
    * @return A completable future to be completed with the listener context.
    */
-  public CompletableFuture<Listener<Long>> onElection(Consumer<Long> listener) {
+  public Listener<Long> onElection(Consumer<Long> callback) {
-    if (!listeners.isEmpty()) {
+    Listener<Long> listener = new Listener<Long>() {
-      listeners.add(listener);
+      @Override
-      return CompletableFuture.completedFuture(new ElectionListener(listener));
+      public void accept(Long epoch) {
-    }
+        callback.accept(epoch);
-
+      }
+      @Override
+      public void close() {
+        listeners.remove(this);
+      }
+    };
     listeners.add(listener);
-    return submit(new LeaderElectionCommands.Listen())
+
-      .thenApply(v -> new ElectionListener(listener));
+    if (epoch != 0) {
+      listener.accept(epoch);
+    }
+    return listener;
   }
 
   /**
@@ -188,33 +198,8 @@ public CompletableFuture<Boolean> isLeader(long epoch) {
    * @return A completable future to be completed once the instance has resigned from leadership for the given epoch.
    */
   public CompletableFuture<Void> resign(long epoch) {
-    return submit(new LeaderElectionCommands.Resign(epoch));
+    return submit(new LeaderElectionCommands.Resign(epoch))
-  }
+      .whenComplete((result, error) -> this.epoch = 0);
-
-  /**
-   * Election listener context.
-   */
-  private class ElectionListener implements Listener<Long> {
-    private final Consumer<Long> listener;
-
-    private ElectionListener(Consumer<Long> listener) {
-      this.listener = listener;
-    }
-
-    @Override
-    public void accept(Long epoch) {
-      listener.accept(epoch);
-    }
-
-    @Override
-    public void close() {
-      synchronized (DistributedLeaderElection.this) {
-        listeners.remove(listener);
-        if (listeners.isEmpty()) {
-          submit(new LeaderElectionCommands.Unlisten());
-        }
-      }
-    }
   }
 
 }
diff --git a/coordination/src/test/java/io/atomix/coordination/DistributedLeaderElectionTest.java b/coordination/src/test/java/io/atomix/coordination/DistributedLeaderElectionTest.java
@@ -42,8 +42,8 @@ public void testElection() throws Throwable {
 
     DistributedLeaderElection election = createResource();
 
-    election.onElection(v -> resume()).thenRun(this::resume);
+    election.onElection(v -> resume());
-    await(0, 2);
+    await(10000);
   }
 
   /**
@@ -60,15 +60,15 @@ public void testNextElection() throws Throwable {
       threadAssertTrue(epoch > lastEpoch.get());
       lastEpoch.set(epoch);
       resume();
-    }).join();
+    });
 
     await(10000);
 
     election2.onElection(epoch -> {
       threadAssertTrue(epoch > lastEpoch.get());
       lastEpoch.set(epoch);
       resume();
-    }).join();
+    });
 
     election1.close();
 
@@ -89,15 +89,15 @@ public void testResignElection() throws Throwable {
       threadAssertTrue(epoch > lastEpoch.get());
       lastEpoch.set(epoch);
       resume();
-    }).join();
+    });
 
     await(10000);
 
     election2.onElection(epoch -> {
       threadAssertTrue(epoch > lastEpoch.get());
       lastEpoch.set(epoch);
       resume();
-    }).join();
+    });
 
     election1.resign(lastEpoch.get());