fix: Fix PayloadProcessor response payload race condition

Motivation An issue was reported in #111 where response payloads logged by a configured PayloadProcessor were sometimes incorrectly empty. Modifications Do not attempt to avoid slicing the response bytebuf in the case that a PayloadProcessor is configured, do not attempt to avoid some additional refcount updates in status != OK case. Result Hopefully no more incorrectly empty logged payloads Resolves #111 Signed-off-by: Nick Hill <[email protected]>
kserve · Nov 23, 2023 · 9ad9bc6 · 9ad9bc6
1 parent a997686
commit 9ad9bc6
Showing 1 changed file with 16 additions and 19 deletions.
diff --git a/src/main/java/com/ibm/watson/modelmesh/ModelMeshApi.java b/src/main/java/com/ibm/watson/modelmesh/ModelMeshApi.java
@@ -725,6 +725,7 @@ public void onHalfClose() {
                 String vModelId = null;
                 String requestId = null;
                 ModelResponse response = null;
+                ByteBuf responsePayload = null;
                 try (InterruptingListener cancelListener = newInterruptingListener()) {
                     if (logHeaders != null) {
                         logHeaders.addToMDC(headers); // MDC cleared in finally block
@@ -767,18 +768,20 @@ public void onHalfClose() {
                     } finally {
                         if (payloadProcessor != null) {
                             processPayload(reqMessage.readerIndex(reqReaderIndex),
-                                    requestId, resolvedModelId, methodName, headers, null, true);
+                                    requestId, resolvedModelId, methodName, headers, null);
                         } else {
                             releaseReqMessage();
                         }
                         reqMessage = null; // ownership released or transferred
                     }
 
-                    respReaderIndex = response.data.readerIndex();
                     respSize = response.data.readableBytes();
                     call.sendHeaders(response.metadata);
+                    if (payloadProcessor != null) {
+                        responsePayload = response.data.retainedSlice();
+                    }
                     call.sendMessage(response.data);
-                    // response is released via ReleaseAfterResponse.releaseAll()
+                    // final response refcount is released via ReleaseAfterResponse.releaseAll()
                     status = OK;
                 } catch (Exception e) {
                     status = toStatus(e);
@@ -795,17 +798,13 @@ public void onHalfClose() {
                         evictMethodDescriptor(methodName);
                     }
                 } finally {
-                    final boolean releaseResponse = status != OK;
                     if (payloadProcessor != null) {
-                        ByteBuf data = null;
-                        Metadata metadata = null;
-                        if (response != null) {
-                            data = response.data.readerIndex(respReaderIndex);
-                            metadata = response.metadata;
-                        }
-                        processPayload(data, requestId, resolvedModelId, methodName, metadata, status, releaseResponse);
-                    } else if (releaseResponse && response != null) {
-                        response.release();
+                        Metadata metadata = response != null ? response.metadata : null;
+                        processPayload(responsePayload, requestId, resolvedModelId, methodName, metadata, status);
+                    }
+                    if (status != OK && response != null) {
+                        // An additional release is required if we call.sendMessage() wasn't sucessful
+                        response.data.release();
                     }
                     ReleaseAfterResponse.releaseAll();
                     clearThreadLocals();
@@ -820,23 +819,21 @@ public void onHalfClose() {
             }
 
             /**
-             * Invoke PayloadProcessor on the request/response data
+             * Invoke PayloadProcessor on the request/response data. This method takes ownership
+             * of the passed-in {@code ByteBuf}.
+             *
              * @param data the binary data
              * @param payloadId the id of the request
              * @param modelId the id of the model
              * @param methodName the name of the invoked method
              * @param metadata the method name metadata
              * @param status null for requests, non-null for responses
-             * @param takeOwnership whether the processor should take ownership
              */
             private void processPayload(ByteBuf data, String payloadId, String modelId, String methodName,
-                                        Metadata metadata, io.grpc.Status status, boolean takeOwnership) {
+                                        Metadata metadata, io.grpc.Status status) {
                 Payload payload = null;
                 try {
                     assert payloadProcessor != null;
-                    if (!takeOwnership) {
-                        ReferenceCountUtil.retain(data);
-                    }
                     payload = new Payload(payloadId, modelId, methodName, metadata, data, status);
                     if (payloadProcessor.process(payload)) {
                         data = null; // ownership transferred