fix(session): fix client keep SessionNotFound if server return ok without response (#220)

* fix(session): fix client keep SessionNotFound if server return ok without response

* move log to callback

* change variable name

* fix the stuck client

* fix scheduler issue

* fix style

Author: mattisonchao

client/src/main/java/io/oxia/client/session/Session.java

@@ -19,6 +19,7 @@
 import static lombok.AccessLevel.PUBLIC;
 
 import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.Throwables;
 import io.grpc.stub.StreamObserver;
 import io.opentelemetry.api.common.Attributes;
 import io.oxia.client.ClientConfig;
@@ -120,7 +121,14 @@ public class Session implements StreamObserver<KeepAliveResponse> {
         this.lastSuccessfullResponse = Instant.now();
         this.heartbeatFuture =
                 executor.scheduleAtFixedRate(
-                        this::sendKeepAlive,
+                        () -> {
+                            // we should catch exception to avoid the schedule future complete
+                            try {
+                                sendKeepAlive();
+                            } catch (Throwable ex) {
+                                log.warn("receive error when send keep-alive request", Throwables.getRootCause(ex));
+                            }
+                        },
                         heartbeatInterval.toMillis(),
                         heartbeatInterval.toMillis(),
                         TimeUnit.MILLISECONDS);
@@ -187,15 +195,7 @@ public CompletableFuture<Void> close() {
                         request,
                         new StreamObserver<>() {
                             @Override
-                            public void onNext(CloseSessionResponse value) {
-                                log.info(
-                                        "Session closed shard={} sessionId={} clientIdentity={}",
-                                        shardId,
-                                        sessionId,
-                                        clientIdentifier);
-                                listener.onSessionClosed(Session.this);
-                                result.complete(null);
-                            }
+                            public void onNext(CloseSessionResponse value) {}
 
                             @Override
                             public void onError(Throwable t) {
@@ -205,9 +205,19 @@ public void onError(Throwable t) {
                             }
 
                             @Override
-                            public void onCompleted() {}
+                            public void onCompleted() {
+                                listener.onSessionClosed(Session.this);
+                                result.complete(null);
+                            }
                         });
 
-        return result;
+        return result.whenComplete(
+                (__, ignore) -> {
+                    log.info(
+                            "Session closed shard={} sessionId={} clientIdentity={}",
+                            shardId,
+                            sessionId,
+                            clientIdentifier);
+                });
     }
 }

client/src/test/java/io/oxia/client/session/SessionTest.java

@@ -168,13 +168,15 @@ public void keepAlive(
             }
 
             responseObserver.onNext(KeepAliveResponse.getDefaultInstance());
+            responseObserver.onCompleted();
         }
 
         @Override
         public void closeSession(
                 CloseSessionRequest request, StreamObserver<CloseSessionResponse> responseObserver) {
             closed.compareAndSet(false, true);
             responseObserver.onNext(CloseSessionResponse.getDefaultInstance());
+            responseObserver.onCompleted();
         }
     }
 }