Refactor DistributedTopic to ensure topic messages are only sent to c…

…lients with registered listeners.

coordination/src/main/java/net/kuujo/copycat/coordination/DistributedTopic.java

@@ -22,9 +22,9 @@
 import net.kuujo.copycat.resource.ResourceContext;
 import net.kuujo.copycat.util.Listener;
 
-import java.util.List;
+import java.util.HashSet;
+import java.util.Set;
 import java.util.concurrent.CompletableFuture;
-import java.util.concurrent.CopyOnWriteArrayList;
 import java.util.function.Consumer;
 
 /**
@@ -33,18 +33,19 @@
  * @author <a href="http://github.com/kuujo">Jordan Halterman</a>
  */
 public class DistributedTopic<T> extends Resource {
-  private final List<TopicListener<T>> listeners = new CopyOnWriteArrayList<>();
+  private final Set<Consumer<T>> listeners = new HashSet<>();
 
   @Override
   protected Class<? extends StateMachine> stateMachine() {
     return TopicState.class;
   }
 
   @Override
+  @SuppressWarnings("unchecked")
   protected void open(ResourceContext context) {
     super.open(context);
     context.session().onReceive(message -> {
-      for (TopicListener<T> listener : listeners) {
+      for (Consumer<T> listener : listeners) {
         listener.accept((T) message);
       }
     });
@@ -68,30 +69,40 @@ public CompletableFuture<Void> publish(T message) {
    * @param listener The message listener.
    * @return The listener context.
    */
-  public Listener<T> onMessage(Consumer<T> listener) {
-    TopicListener<T> context = new TopicListener<T>(listener);
-    listeners.add(context);
-    return context;
+  public CompletableFuture<Listener<T>> onMessage(Consumer<T> listener) {
+    if (!listeners.isEmpty()) {
+      listeners.add(listener);
+      return CompletableFuture.completedFuture(new TopicListener(listener));
+    }
+
+    listeners.add(listener);
+    return submit(TopicCommands.Listen.builder().build())
+      .thenApply(v -> new TopicListener(listener));
   }
 
   /**
-   * Topic listener context.
+   * Topic listener.
    */
-  private class TopicListener<T> implements Listener<T> {
+  private class TopicListener implements Listener<T> {
     private final Consumer<T> listener;
 
     private TopicListener(Consumer<T> listener) {
       this.listener = listener;
     }
 
     @Override
-    public void accept(T event) {
-      listener.accept(event);
+    public void accept(T message) {
+      listener.accept(message);
     }
 
     @Override
     public void close() {
-      listeners.remove(this);
+      synchronized (DistributedTopic.this) {
+        listeners.remove(listener);
+        if (listeners.isEmpty()) {
+          submit(TopicCommands.Unlisten.builder().build());
+        }
+      }
     }
   }
 

coordination/src/main/java/net/kuujo/copycat/coordination/state/TopicCommands.java

@@ -15,13 +15,14 @@
  */
 package net.kuujo.copycat.coordination.state;
 
-import net.kuujo.copycat.util.BuilderPool;
-import net.kuujo.copycat.raft.protocol.Command;
-import net.kuujo.copycat.raft.protocol.Operation;
 import net.kuujo.copycat.io.BufferInput;
 import net.kuujo.copycat.io.BufferOutput;
 import net.kuujo.copycat.io.serializer.CopycatSerializable;
+import net.kuujo.copycat.io.serializer.SerializeWith;
 import net.kuujo.copycat.io.serializer.Serializer;
+import net.kuujo.copycat.raft.protocol.Command;
+import net.kuujo.copycat.raft.protocol.Operation;
+import net.kuujo.copycat.util.BuilderPool;
 
 /**
  * Topic commands.
@@ -37,6 +38,15 @@ private TopicCommands() {
    * Abstract topic command.
    */
   public static abstract class TopicCommand<V> implements Command<V>, CopycatSerializable {
+    @Override
+    public void writeObject(BufferOutput buffer, Serializer serializer) {
+
+    }
+
+    @Override
+    public void readObject(BufferInput buffer, Serializer serializer) {
+
+    }
 
     /**
      * Base map command builder.
@@ -48,6 +58,68 @@ protected Builder(BuilderPool<T, U> pool) {
     }
   }
 
+  /**
+   * Listen command.
+   */
+  @SerializeWith(id=510)
+  public static class Listen extends TopicCommand<Void> {
+
+    /**
+     * Returns a new listen command builder.
+     *
+     * @return A new listen command builder.
+     */
+    @SuppressWarnings("unchecked")
+    public static Builder builder() {
+      return Operation.builder(Builder.class, Builder::new);
+    }
+
+    /**
+     * Listen command builder.
+     */
+    public static class Builder extends TopicCommand.Builder<Builder, Listen, Void> {
+      public Builder(BuilderPool<Builder, Listen> pool) {
+        super(pool);
+      }
+
+      @Override
+      protected Listen create() {
+        return new Listen();
+      }
+    }
+  }
+
+  /**
+   * Unlisten command.
+   */
+  @SerializeWith(id=511)
+  public static class Unlisten extends TopicCommand<Void> {
+
+    /**
+     * Returns a new unlisten command builder.
+     *
+     * @return A new unlisten command builder.
+     */
+    @SuppressWarnings("unchecked")
+    public static Builder builder() {
+      return Operation.builder(Builder.class, Builder::new);
+    }
+
+    /**
+     * Unlisten command builder.
+     */
+    public static class Builder extends TopicCommand.Builder<Builder, Unlisten, Void> {
+      public Builder(BuilderPool<Builder, Unlisten> pool) {
+        super(pool);
+      }
+
+      @Override
+      protected Unlisten create() {
+        return new Unlisten();
+      }
+    }
+  }
+
   /**
    * Publish command.
    */
@@ -112,44 +184,4 @@ protected Publish<T> create() {
     }
   }
 
-  /**
-   * Subscribe command.
-   */
-  public static class Subscribe extends TopicCommand<Void> {
-
-    /**
-     * Returns a new publish command builder.
-     *
-     * @return The publish command builder.
-     */
-    @SuppressWarnings("unchecked")
-    public static Builder builder() {
-      return Operation.builder(Builder.class, Builder::new);
-    }
-
-    @Override
-    public void writeObject(BufferOutput buffer, Serializer serializer) {
-
-    }
-
-    @Override
-    public void readObject(BufferInput buffer, Serializer serializer) {
-
-    }
-
-    /**
-     * Publish command builder.
-     */
-    public static class Builder extends TopicCommand.Builder<Builder, Subscribe, Void> {
-      public Builder(BuilderPool<Builder, Subscribe> pool) {
-        super(pool);
-      }
-
-      @Override
-      protected Subscribe create() {
-        return new Subscribe();
-      }
-    }
-  }
-
 }

coordination/src/main/java/net/kuujo/copycat/coordination/state/TopicState.java

@@ -15,29 +15,55 @@
  */
 package net.kuujo.copycat.coordination.state;
 
-import net.kuujo.copycat.raft.session.Session;
 import net.kuujo.copycat.raft.Commit;
 import net.kuujo.copycat.raft.StateMachine;
 import net.kuujo.copycat.raft.StateMachineExecutor;
 
+import java.util.HashMap;
+import java.util.Map;
+
 /**
  * Topic state machine.
  *
  * @author <a href="http://github.com/kuujo">Jordan Halterman</a>
  */
 public class TopicState extends StateMachine {
+  private final Map<Long, Commit<TopicCommands.Listen>> listeners = new HashMap<>();
 
   @Override
   public void configure(StateMachineExecutor executor) {
     executor.register(TopicCommands.Publish.class, this::publish);
   }
 
+  /**
+   * Applies listen commits.
+   */
+  protected void listen(Commit<TopicCommands.Listen> commit) {
+    if (!listeners.containsKey(commit.session().id())) {
+      listeners.put(commit.session().id(), commit);
+    } else {
+      commit.clean();
+    }
+  }
+
+  /**
+   * Applies listen commits.
+   */
+  protected void unlisten(Commit<LeaderElectionCommands.Unlisten> commit) {
+    Commit<TopicCommands.Listen> listener = listeners.remove(commit.session().id());
+    if (listener != null) {
+      listener.clean();
+    } else {
+      commit.clean();
+    }
+  }
+
   /**
    * Handles a publish commit.
    */
   protected void publish(Commit<TopicCommands.Publish> commit) {
-    for (Session session : context().sessions()) {
-      session.publish(commit.operation().message());
+    for (Commit<TopicCommands.Listen> listener : listeners.values()) {
+      listener.session().publish(commit.operation().message());
     }
     commit.clean();
   }