feat(proxy): optimize SuspendRequestService (#585)

Signed-off-by: SSpirits <[email protected]>
AutoMQ · Nov 6, 2023 · 4f1df04 · 4f1df04
1 parent 30d0e3e
commit 4f1df04
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 2 deletions.
diff --git a/.../com/automq/rocketmq/proxy/grpc/v2/consumer/ExtendReceiveMessageResponseStreamWriter.java b/.../com/automq/rocketmq/proxy/grpc/v2/consumer/ExtendReceiveMessageResponseStreamWriter.java
@@ -52,6 +52,7 @@ public void writeAndComplete(ProxyContext ctx, ReceiveMessageRequest request, Po
                 } else {
                     recordRpcLatency(ctx, Code.OK);
                 }
+                break;
             case POLLING_FULL:
                 recordRpcLatency(ctx, Code.TOO_MANY_REQUESTS);
                 break;

diff --git a/proxy/src/main/java/com/automq/rocketmq/proxy/service/SuspendRequestService.java b/proxy/src/main/java/com/automq/rocketmq/proxy/service/SuspendRequestService.java
@@ -143,9 +143,15 @@ public CompletableFuture<Boolean> tryFetchMessages() {
                             completed.set(true);
                             return true;
                         }
-                        inflight.set(false);
                         return false;
-                    });
+                    })
+                    .exceptionally(ex -> {
+                        LOGGER.error("Error while fetching messages for suspended request.", ex);
+                        future.completeExceptionally(ex);
+                        completed.set(true);
+                        return true;
+                    })
+                    .whenComplete((result, ex) -> inflight.set(false));
             }
             return CompletableFuture.completedFuture(false);
         }
@@ -171,6 +177,7 @@ public void notifyMessageArrival(String topic, int queueId, String tag) {
                                 suspendRequestCount.decrementAndGet();
                             }
                         }));
+                break;
             }
         }
     }