Fixing memory leak

Author: FabianMeiswinkel

sdk/cosmos/azure-cosmos-encryption/src/test/java/com/azure/cosmos/encryption/CosmosNettyLeakDetectorFactory.java

@@ -54,6 +54,7 @@ public static void ingestIntoNetty() {
             ResourceLeakDetector.setLevel(ResourceLeakDetector.Level.PARANOID);
             // sample every allocation
             System.setProperty("io.netty.leakDetection.samplingInterval", "1");
+            System.setProperty("io.netty.leakDetection.targetRecords", "256");
             // install custom reporter
             ResourceLeakDetectorFactory.setResourceLeakDetectorFactory(new CosmosNettyLeakDetectorFactory());
 

sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosNettyLeakDetectorFactory.java

@@ -54,6 +54,7 @@ public static void ingestIntoNetty() {
             ResourceLeakDetector.setLevel(ResourceLeakDetector.Level.PARANOID);
             // sample every allocation
             System.setProperty("io.netty.leakDetection.samplingInterval", "1");
+            System.setProperty("io.netty.leakDetection.targetRecords", "256");
             // install custom reporter
             ResourceLeakDetectorFactory.setResourceLeakDetectorFactory(new CosmosNettyLeakDetectorFactory());
 

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxGatewayStoreModel.java

@@ -417,46 +417,60 @@ private Mono<RxDocumentServiceResponse> toDocumentServiceResponse(Mono<HttpRespo
                 .map(bodyByteBuf -> leakDetectionDebuggingEnabled
                     ? bodyByteBuf.retain().touch(this)
                     : bodyByteBuf.retain())
-                .publishOn(CosmosSchedulers.TRANSPORT_RESPONSE_BOUNDED_ELASTIC);
+                .publishOn(CosmosSchedulers.TRANSPORT_RESPONSE_BOUNDED_ELASTIC)
+                .doOnDiscard(ByteBuf.class, buf -> {
+                    if (buf.refCnt() > 0) {
+                        io.netty.util.ReferenceCountUtil.safeRelease(buf);
+                    }
+                });
 
             return contentObservable
                 .map(content -> {
                     if (leakDetectionDebuggingEnabled) {
                         content.touch(this);
                     }
 
-                    // Capture transport client request timeline
-                    ReactorNettyRequestRecord reactorNettyRequestRecord = httpResponse.request().reactorNettyRequestRecord();
-                    if (reactorNettyRequestRecord != null) {
-                        reactorNettyRequestRecord.setTimeCompleted(Instant.now());
-                    }
+                    try {
+                        // Capture transport client request timeline
+                        ReactorNettyRequestRecord reactorNettyRequestRecord = httpResponse.request().reactorNettyRequestRecord();
+                        if (reactorNettyRequestRecord != null) {
+                            reactorNettyRequestRecord.setTimeCompleted(Instant.now());
+                        }
 
-                    StoreResponse rsp = request
-                        .getEffectiveHttpTransportSerializer(this)
-                        .unwrapToStoreResponse(httpRequest.uri().toString(), request, httpResponseStatus, httpResponseHeaders, content);
+                        StoreResponse rsp = request
+                            .getEffectiveHttpTransportSerializer(this)
+                            .unwrapToStoreResponse(httpRequest.uri().toString(), request, httpResponseStatus, httpResponseHeaders, content);
 
-                    if (reactorNettyRequestRecord != null) {
-                        rsp.setRequestTimeline(reactorNettyRequestRecord.takeTimelineSnapshot());
+                        if (reactorNettyRequestRecord != null) {
+                            rsp.setRequestTimeline(reactorNettyRequestRecord.takeTimelineSnapshot());
 
-                        if (this.gatewayServerErrorInjector != null) {
-                            // only configure when fault injection is used
-                            rsp.setFaultInjectionRuleId(
-                                request
-                                    .faultInjectionRequestContext
-                                    .getFaultInjectionRuleId(reactorNettyRequestRecord.getTransportRequestId()));
+                            if (this.gatewayServerErrorInjector != null) {
+                                // only configure when fault injection is used
+                                rsp.setFaultInjectionRuleId(
+                                    request
+                                        .faultInjectionRequestContext
+                                        .getFaultInjectionRuleId(reactorNettyRequestRecord.getTransportRequestId()));
 
-                            rsp.setFaultInjectionRuleEvaluationResults(
-                                request
-                                    .faultInjectionRequestContext
-                                    .getFaultInjectionRuleEvaluationResults(reactorNettyRequestRecord.getTransportRequestId()));
+                                rsp.setFaultInjectionRuleEvaluationResults(
+                                    request
+                                        .faultInjectionRequestContext
+                                        .getFaultInjectionRuleEvaluationResults(reactorNettyRequestRecord.getTransportRequestId()));
+                            }
                         }
-                    }
 
-                    if (request.requestContext.cosmosDiagnostics != null) {
-                        BridgeInternal.recordGatewayResponse(request.requestContext.cosmosDiagnostics, request, rsp, globalEndpointManager);
-                    }
+                        if (request.requestContext.cosmosDiagnostics != null) {
+                            BridgeInternal.recordGatewayResponse(request.requestContext.cosmosDiagnostics, request, rsp, globalEndpointManager);
+                        }
 
-                    return rsp;
+                        return rsp;
+                    } catch (Throwable t) {
+                        if (content.refCnt() > 0) {
+                            // Unwrap failed before StoreResponse took ownership -> release our retain
+                            io.netty.util.ReferenceCountUtil.safeRelease(content);
+                        }
+
+                        throw t;
+                    }
                 })
                 .single();