fixes memory leak when default rsocket tracing propagation is enabled (#2262)

* fixes memory leak when default rsocket tracing propagation is enabled

Signed-off-by: OlegDokuka <[email protected]>

* fixes styles

Signed-off-by: OlegDokuka <[email protected]>

---------

Signed-off-by: OlegDokuka <[email protected]>

Author: OlegDokuka

spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/rsocket/TracingRequesterRSocketProxy.java

@@ -21,10 +21,12 @@
 import java.util.function.Function;
 
 import io.netty.buffer.ByteBuf;
+import io.netty.buffer.ByteBufAllocator;
 import io.netty.buffer.CompositeByteBuf;
 import io.rsocket.Payload;
 import io.rsocket.RSocket;
 import io.rsocket.frame.FrameType;
+import io.rsocket.metadata.CompositeMetadataCodec;
 import io.rsocket.metadata.RoutingMetadata;
 import io.rsocket.metadata.TracingMetadataCodec;
 import io.rsocket.metadata.WellKnownMimeType;
@@ -104,34 +106,40 @@ <T> Mono<T> setSpan(Function<Payload, Mono<T>> input, Payload payload, FrameType
 				log.debug("Extracted result from context or thread local " + span);
 			}
 			final Payload newPayload = PayloadUtils.cleanTracingMetadata(payload, new HashSet<>(propagator.fields()));
-			TraceContext traceContext = span.context();
+			final TraceContext traceContext = span.context();
+			final CompositeByteBuf metadata = (CompositeByteBuf) newPayload.metadata();
 			if (this.isZipkinPropagationEnabled) {
-				injectDefaultZipkinRSocketHeaders(newPayload, traceContext);
+				injectDefaultZipkinRSocketHeaders(metadata, traceContext);
 			}
-			this.propagator.inject(traceContext, (CompositeByteBuf) newPayload.metadata(), this.setter);
+			this.propagator.inject(traceContext, metadata, this.setter);
 			return input.apply(newPayload).doOnError(span::error).doFinally(signalType -> span.end());
 		});
 	}
 
-	private void injectDefaultZipkinRSocketHeaders(Payload newPayload, TraceContext traceContext) {
+	void injectDefaultZipkinRSocketHeaders(CompositeByteBuf metadata, TraceContext traceContext) {
 		TracingMetadataCodec.Flags flags = traceContext.sampled() == null ? TracingMetadataCodec.Flags.UNDECIDED
 				: traceContext.sampled() ? TracingMetadataCodec.Flags.SAMPLE : TracingMetadataCodec.Flags.NOT_SAMPLE;
 		String traceId = traceContext.traceId();
 		long[] traceIds = EncodingUtils.fromString(traceId);
 		long[] spanId = EncodingUtils.fromString(traceContext.spanId());
 		long[] parentSpanId = EncodingUtils.fromString(traceContext.parentId());
 		boolean isTraceId128Bit = traceIds.length == 2;
+
+		final ByteBufAllocator allocator = metadata.alloc();
 		if (isTraceId128Bit) {
-			TracingMetadataCodec.encode128(newPayload.metadata().alloc(), traceIds[0], traceIds[1], spanId[0],
-					EncodingUtils.fromString(traceContext.parentId())[0], flags);
+			CompositeMetadataCodec.encodeAndAddMetadata(metadata, allocator,
+					WellKnownMimeType.MESSAGE_RSOCKET_TRACING_ZIPKIN,
+					TracingMetadataCodec.encode128(allocator, traceIds[0], traceIds[1], spanId[0],
+							EncodingUtils.fromString(traceContext.parentId())[0], flags));
 		}
 		else {
-			TracingMetadataCodec.encode64(newPayload.metadata().alloc(), traceIds[0], spanId[0], parentSpanId[0],
-					flags);
+			CompositeMetadataCodec.encodeAndAddMetadata(metadata, allocator,
+					WellKnownMimeType.MESSAGE_RSOCKET_TRACING_ZIPKIN,
+					TracingMetadataCodec.encode64(allocator, traceIds[0], spanId[0], parentSpanId[0], flags));
 		}
 	}
 
-	private Span.Builder spanBuilder(ContextView contextView) {
+	Span.Builder spanBuilder(ContextView contextView) {
 		Span.Builder spanBuilder = this.tracer.spanBuilder();
 		if (contextView.hasKey(TraceContext.class)) {
 			spanBuilder = spanBuilder.setParent(contextView.get(TraceContext.class));

spring-cloud-sleuth-instrumentation/src/test/java/org/springframework/cloud/sleuth/instrument/rsocket/TracingRequesterRSocketProxyTest.java

@@ -0,0 +1,74 @@
+/*
+ * Copyright 2013-2023 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.cloud.sleuth.instrument.rsocket;
+
+import io.netty.buffer.ByteBuf;
+import io.netty.buffer.ByteBufAllocator;
+import io.netty.buffer.CompositeByteBuf;
+import org.assertj.core.api.Assertions;
+import org.junit.jupiter.api.Test;
+
+import org.springframework.cloud.sleuth.TraceContext;
+
+public class TracingRequesterRSocketProxyTest {
+
+	@Test
+	public void checkNoLeaksOnDefaultTracingHeaders() {
+		final TracingRequesterRSocketProxy tracingRequesterRSocketProxy = new TracingRequesterRSocketProxy(null, null,
+				null, null, true);
+
+		final ByteBufAllocator allocator = ByteBufAllocator.DEFAULT;
+		final CompositeByteBuf metadata = allocator.compositeBuffer();
+		tracingRequesterRSocketProxy.injectDefaultZipkinRSocketHeaders(metadata, new TraceContext() {
+			@Override
+			public String traceId() {
+				return "0000000000000002";
+			}
+
+			@Override
+			public String parentId() {
+				return null;
+			}
+
+			@Override
+			public String spanId() {
+				return "0000000000000001";
+			}
+
+			@Override
+			public Boolean sampled() {
+				return null;
+			}
+		});
+
+		// should add 2 components which are headers and body
+		Assertions.assertThat(metadata.numComponents()).isEqualTo(2);
+		Assertions.assertThat(metadata.refCnt()).isEqualTo(1);
+
+		final ByteBuf c1 = metadata.internalComponent(0);
+		final ByteBuf c2 = metadata.internalComponent(1);
+
+		Assertions.assertThat(c1.refCnt()).isEqualTo(1);
+		Assertions.assertThat(c2.refCnt()).isEqualTo(1);
+
+		Assertions.assertThat(metadata.release()).isTrue();
+
+		Assertions.assertThat(c1.refCnt()).isEqualTo(0);
+		Assertions.assertThat(c2.refCnt()).isEqualTo(0);
+	}
+
+}