Update for review comments

Signed-off-by: Andrew Carbonetto <[email protected]>

java/.gitignore

@@ -5,4 +5,4 @@
 build
 
 # Ignore protobuf generated files
-models/protobuf
+protobuf

java/client/src/main/java/babushka/connectors/SocketConnection.java

@@ -117,10 +117,7 @@ public void close() {
     channel.close();
     group.shutdownGracefully();
     INSTANCE = null;
-    // TODO should we reply in uncompleted futures?
     CallbackManager.shutdownGracefully();
-    CallbackManager.connectionRequests.clear();
-    CallbackManager.responses.clear();
   }
 
   /**

java/client/src/main/java/babushka/connectors/handlers/ChannelBuilder.java

@@ -16,7 +16,7 @@ public void initChannel(@NonNull UnixChannel ch) {
         // https://netty.io/4.1/api/io/netty/handler/codec/protobuf/ProtobufEncoder.html
         .addLast("protobufDecoder", new ProtobufVarint32FrameDecoder())
         .addLast("protobufEncoder", new ProtobufVarint32LengthFieldPrepender())
-        .addLast(new ChannelInboundHandlerAdapter())
-        .addLast(new ChannelOutboundHandlerAdapter());
+        .addLast(new ReadHandler())
+        .addLast(new WriteHandler());
   }
 }

java/client/src/main/java/babushka/connectors/handlers/ReadHandler.java

@@ -26,16 +26,7 @@ public void channelRead(@NonNull ChannelHandlerContext ctx, @NonNull Object msg)
     buf.readBytes(bytes);
     // TODO surround parsing with try-catch, set error to future if parsing failed.
     var response = ResponseOuterClass.Response.parseFrom(bytes);
-    int callbackId = response.getCallbackIdx();
-    if (callbackId == 0) {
-      // can't distinguish connection requests since they have no
-      // callback ID
-      // https://github.com/aws/babushka/issues/600
-      CallbackManager.connectionRequests.pop().complete(response);
-    } else {
-      CallbackManager.responses.get(callbackId).complete(response);
-      CallbackManager.responses.remove(callbackId);
-    }
+    CallbackManager.completeAsync(response);
     buf.release();
   }
 

java/client/src/main/java/babushka/managers/CallbackManager.java

@@ -58,15 +58,30 @@ public void registerConnection(CompletableFuture<Response> future) {
     CallbackManager.connectionRequests.add(future);
   }
 
+  public static void completeAsync(Response response) {
+    int callbackId = response.getCallbackIdx();
+    if (callbackId == 0) {
+      // can't distinguish connection requests since they have no
+      // callback ID
+      // https://github.com/aws/babushka/issues/600
+      connectionRequests.pop().completeAsync(() -> response);
+    } else {
+      responses.get(callbackId).completeAsync(() -> response);
+      responses.remove(callbackId);
+    }
+  }
+
   public static void shutdownGracefully() {
     connectionRequests.forEach(
         future -> {
           future.completeExceptionally(new InterruptedException());
         });
+    connectionRequests.clear();
     responses.forEach(
         (callbackId, future) -> {
           future.completeExceptionally(new InterruptedException());
         });
+    responses.clear();
   }
 
   /**

java/client/src/main/java/babushka/managers/CommandManager.java

@@ -50,14 +50,14 @@ public CompletableFuture<String> set(String key, String value) {
   public CompletableFuture<String> submitNewCommand(RequestType command, List<String> args) {
     // TODO this explicitly uses ForkJoin thread pool. May be we should use another one.
     CompletableFuture<Response> future = new CompletableFuture<>();
-    int callbackId = callbackManager.registerRequest(future);
 
     return CompletableFuture.supplyAsync(
             () -> {
               socketConnection.writeAndFlush(redisSingleCommand(command, args));
               return future;
             })
-        .thenCompose(f -> f)
-        .thenApply(RequestBuilder::resolveRedisResponseToString);
+        // TODO: is there a better way to execute this?
+        .thenComposeAsync(f -> f)
+        .thenApplyAsync(RequestBuilder::resolveRedisResponseToString);
   }
 }

java/client/src/main/java/babushka/managers/ConnectionManager.java

@@ -52,7 +52,7 @@ public CompletableFuture<String> connectToRedis(
     var future = new CompletableFuture<Response>();
     callbackManager.registerConnection(future);
     socketConnection.writeAndFlush(request);
-    return future.thenApply(f -> f.getConstantResponse().toString());
+    return future.thenApplyAsync(f -> f.getConstantResponse().toString());
   }
 
   /**