get(long streamId, long startOffset, long endOffset, int maxBytes) {
+ return get(TraceContext.DEFAULT, streamId, startOffset, endOffset, maxBytes);
+ }
+
/**
* Get streamId [startOffset, endOffset) range records with maxBytes limit.
*
@@ -124,7 +131,13 @@ public boolean put(StreamRecordBatch recordBatch) {
*
* Note: the records is retained, the caller should release it.
*/
- public List get(long streamId, long startOffset, long endOffset, int maxBytes) {
+ @WithSpan
+ public List get(TraceContext context,
+ @SpanAttribute long streamId,
+ @SpanAttribute long startOffset,
+ @SpanAttribute long endOffset,
+ @SpanAttribute int maxBytes) {
+ context.currentContext();
TimerUtil timerUtil = new TimerUtil();
List records;
readLock.lock();
diff --git a/s3stream/src/main/java/com/automq/stream/s3/cache/S3BlockCache.java b/s3stream/src/main/java/com/automq/stream/s3/cache/S3BlockCache.java
index 060cc7600..161780a29 100644
--- a/s3stream/src/main/java/com/automq/stream/s3/cache/S3BlockCache.java
+++ b/s3stream/src/main/java/com/automq/stream/s3/cache/S3BlockCache.java
@@ -17,6 +17,8 @@
package com.automq.stream.s3.cache;
+import com.automq.stream.s3.trace.context.TraceContext;
+
import java.util.concurrent.CompletableFuture;
/**
@@ -26,5 +28,9 @@
*/
public interface S3BlockCache {
- CompletableFuture read(long streamId, long startOffset, long endOffset, int maxBytes);
+ CompletableFuture read(TraceContext context, long streamId, long startOffset, long endOffset, int maxBytes);
+
+ default CompletableFuture read(long streamId, long startOffset, long endOffset, int maxBytes) {
+ return read(TraceContext.DEFAULT, streamId, startOffset, endOffset, maxBytes);
+ }
}
diff --git a/s3stream/src/main/java/com/automq/stream/s3/cache/StreamReader.java b/s3stream/src/main/java/com/automq/stream/s3/cache/StreamReader.java
index 06640b21e..8981fdc83 100644
--- a/s3stream/src/main/java/com/automq/stream/s3/cache/StreamReader.java
+++ b/s3stream/src/main/java/com/automq/stream/s3/cache/StreamReader.java
@@ -26,7 +26,12 @@
import com.automq.stream.s3.model.StreamRecordBatch;
import com.automq.stream.s3.objects.ObjectManager;
import com.automq.stream.s3.operator.S3Operator;
+import com.automq.stream.s3.trace.TraceUtils;
+import com.automq.stream.s3.trace.context.TraceContext;
import com.automq.stream.utils.Threads;
+import io.opentelemetry.api.common.Attributes;
+import io.opentelemetry.instrumentation.annotations.SpanAttribute;
+import io.opentelemetry.instrumentation.annotations.WithSpan;
import org.apache.commons.lang3.tuple.ImmutablePair;
import org.apache.commons.lang3.tuple.Pair;
import org.slf4j.Logger;
@@ -105,8 +110,12 @@ public void shutdown() {
errorHandlerExecutor.shutdown();
}
- public CompletableFuture> syncReadAhead(long streamId, long startOffset, long endOffset,
- int maxBytes, ReadAheadAgent agent, UUID uuid) {
+ @WithSpan
+ public CompletableFuture> syncReadAhead(TraceContext traceContext,
+ @SpanAttribute long streamId,
+ @SpanAttribute long startOffset,
+ @SpanAttribute long endOffset,
+ @SpanAttribute int maxBytes, ReadAheadAgent agent, UUID uuid) {
if (LOGGER.isDebugEnabled()) {
LOGGER.debug("[S3BlockCache] sync read ahead, stream={}, {}-{}, maxBytes={}", streamId, startOffset, endOffset, maxBytes);
}
@@ -119,9 +128,9 @@ public CompletableFuture> syncReadAhead(long streamId, l
if (inflightReadAheadTaskMap.putIfAbsent(readAheadTaskKey, readAheadTaskContext) == null) {
context.taskKeySet.add(readAheadTaskKey);
}
- return getDataBlockIndices(streamId, endOffset, context)
+ return getDataBlockIndices(traceContext, streamId, endOffset, context)
.thenComposeAsync(v ->
- handleSyncReadAhead(streamId, startOffset, endOffset, maxBytes, agent, uuid, timer, context), streamReaderExecutor)
+ handleSyncReadAhead(traceContext, streamId, startOffset, endOffset, maxBytes, agent, uuid, timer, context), streamReaderExecutor)
.whenComplete((nil, ex) -> {
for (DefaultS3BlockCache.ReadAheadTaskKey key : context.taskKeySet) {
completeInflightTask0(key, ex);
@@ -131,7 +140,8 @@ public CompletableFuture> syncReadAhead(long streamId, l
});
}
- CompletableFuture> handleSyncReadAhead(long streamId, long startOffset, long endOffset,
+ @WithSpan
+ CompletableFuture> handleSyncReadAhead(TraceContext traceContext, long streamId, long startOffset, long endOffset,
int maxBytes, ReadAheadAgent agent, UUID uuid,
TimerUtil timer, ReadContext context) {
if (context.streamDataBlocksPair.isEmpty()) {
@@ -158,7 +168,7 @@ CompletableFuture> handleSyncReadAhead(long streamId, lo
totalReserveSize, uuid, streamId, startOffset, endOffset, maxBytes);
}
- CompletableFuture throttleCf = inflightReadThrottle.acquire(uuid, totalReserveSize);
+ CompletableFuture throttleCf = inflightReadThrottle.acquire(traceContext, uuid, totalReserveSize);
return throttleCf.thenComposeAsync(nil -> {
// concurrently read all data blocks
for (int i = 0; i < streamDataBlocksToRead.size(); i++) {
@@ -182,31 +192,37 @@ CompletableFuture> handleSyncReadAhead(long streamId, lo
setInflightReadAheadStatus(new DefaultS3BlockCache.ReadAheadTaskKey(streamId, startOffset),
DefaultS3BlockCache.ReadBlockCacheStatus.WAIT_FETCH_DATA);
}
- cfList.add(reserveResult.cf().thenApplyAsync(dataBlock -> {
- if (dataBlock.records().isEmpty()) {
- return new ArrayList();
- }
- // retain records to be returned
- dataBlock.records().forEach(StreamRecordBatch::retain);
- recordsMap.put(dataBlockKey, dataBlock.records());
-
- // retain records to be put into block cache
- dataBlock.records().forEach(StreamRecordBatch::retain);
- blockCache.put(streamId, dataBlock.records());
- dataBlock.release();
-
- return dataBlock.records();
- }, backgroundExecutor).whenComplete((ret, ex) -> {
- if (ex != null) {
- LOGGER.error("[S3BlockCache] sync ra fail to read data block, stream={}, {}-{}, data block: {}",
- streamId, startOffset, endOffset, streamDataBlock, ex);
- }
- completeInflightTask(context, taskKey, ex);
- if (isNotAlignedFirstBlock) {
- // in case of first data block and startOffset is not aligned with start of data block
- completeInflightTask(context, new DefaultS3BlockCache.ReadAheadTaskKey(streamId, startOffset), ex);
- }
- }));
+ try {
+ CompletableFuture> cf = TraceUtils.runWithSpanAsync(new TraceContext(traceContext), Attributes.empty(), "StreamReader::readDataBlock",
+ () -> reserveResult.cf().thenApplyAsync(dataBlock -> {
+ if (dataBlock.records().isEmpty()) {
+ return new ArrayList();
+ }
+ // retain records to be returned
+ dataBlock.records().forEach(StreamRecordBatch::retain);
+ recordsMap.put(dataBlockKey, dataBlock.records());
+
+ // retain records to be put into block cache
+ dataBlock.records().forEach(StreamRecordBatch::retain);
+ blockCache.put(streamId, dataBlock.records());
+ dataBlock.release();
+
+ return dataBlock.records();
+ }, backgroundExecutor).whenComplete((ret, ex) -> {
+ if (ex != null) {
+ LOGGER.error("[S3BlockCache] sync ra fail to read data block, stream={}, {}-{}, data block: {}",
+ streamId, startOffset, endOffset, streamDataBlock, ex);
+ }
+ completeInflightTask(context, taskKey, ex);
+ if (isNotAlignedFirstBlock) {
+ // in case of first data block and startOffset is not aligned with start of data block
+ completeInflightTask(context, new DefaultS3BlockCache.ReadAheadTaskKey(streamId, startOffset), ex);
+ }
+ }));
+ cfList.add(cf);
+ } catch (Throwable e) {
+ throw new IllegalArgumentException(e);
+ }
if (reserveResult.reserveSize() > 0) {
dataBlockReadAccumulator.readDataBlock(objectReader, streamDataBlock.dataBlockIndex());
}
@@ -274,7 +290,7 @@ public void asyncReadAhead(long streamId, long startOffset, long endOffset, int
DefaultS3BlockCache.ReadBlockCacheStatus.WAIT_DATA_INDEX);
inflightReadAheadTaskMap.putIfAbsent(readAheadTaskKey, readAheadTaskContext);
context.taskKeySet.add(readAheadTaskKey);
- getDataBlockIndices(streamId, endOffset, context)
+ getDataBlockIndices(TraceContext.DEFAULT, streamId, endOffset, context)
.thenAcceptAsync(v ->
handleAsyncReadAhead(streamId, startOffset, endOffset, maxBytes, agent, timer, context), streamReaderExecutor)
.whenComplete((nil, ex) -> {
@@ -347,7 +363,7 @@ CompletableFuture handleAsyncReadAhead(long streamId, long startOffset, lo
reserveResult.reserveSize(), uuid, streamId, startOffset, endOffset, maxBytes);
}
if (reserveResult.reserveSize() > 0) {
- inflightReadThrottle.acquire(uuid, reserveResult.reserveSize()).thenAcceptAsync(nil -> {
+ inflightReadThrottle.acquire(TraceContext.DEFAULT, uuid, reserveResult.reserveSize()).thenAcceptAsync(nil -> {
// read data block
if (context.taskKeySet.contains(taskKey)) {
setInflightReadAheadStatus(taskKey, DefaultS3BlockCache.ReadBlockCacheStatus.WAIT_FETCH_DATA);
@@ -380,7 +396,9 @@ CompletableFuture handleAsyncReadAhead(long streamId, long startOffset, lo
});
}
- CompletableFuture getDataBlockIndices(long streamId, long endOffset, ReadContext context) {
+ @WithSpan
+ CompletableFuture getDataBlockIndices(TraceContext traceContext, long streamId, long endOffset, ReadContext context) {
+ traceContext.currentContext();
CompletableFuture getObjectsCf = CompletableFuture.completedFuture(false);
if (context.objectIndex >= context.objects.size()) {
getObjectsCf = objectManager
@@ -450,7 +468,7 @@ CompletableFuture getDataBlockIndices(long streamId, long endOffset, ReadC
if (findIndexResult.isFulfilled()) {
return CompletableFuture.completedFuture(null);
}
- return getDataBlockIndices(streamId, endOffset, context);
+ return getDataBlockIndices(traceContext, streamId, endOffset, context);
}, streamReaderExecutor);
}
diff --git a/s3stream/src/main/java/com/automq/stream/s3/context/AppendContext.java b/s3stream/src/main/java/com/automq/stream/s3/context/AppendContext.java
new file mode 100644
index 000000000..de8ef8626
--- /dev/null
+++ b/s3stream/src/main/java/com/automq/stream/s3/context/AppendContext.java
@@ -0,0 +1,38 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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
+ *
+ * http://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 com.automq.stream.s3.context;
+
+import com.automq.stream.s3.trace.context.TraceContext;
+import io.opentelemetry.api.trace.Tracer;
+import io.opentelemetry.context.Context;
+
+public class AppendContext extends TraceContext {
+ public static final AppendContext DEFAULT = new AppendContext();
+
+ public AppendContext() {
+ super(false, null, null);
+ }
+
+ public AppendContext(TraceContext context) {
+ super(context);
+ }
+
+ public AppendContext(boolean isTraceEnabled, Tracer tracer, Context currentContext) {
+ super(isTraceEnabled, tracer, currentContext);
+ }
+}
diff --git a/s3stream/src/main/java/com/automq/stream/s3/context/FetchContext.java b/s3stream/src/main/java/com/automq/stream/s3/context/FetchContext.java
new file mode 100644
index 000000000..3218e864b
--- /dev/null
+++ b/s3stream/src/main/java/com/automq/stream/s3/context/FetchContext.java
@@ -0,0 +1,48 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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
+ *
+ * http://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 com.automq.stream.s3.context;
+
+import com.automq.stream.api.ReadOptions;
+import com.automq.stream.s3.trace.context.TraceContext;
+import io.opentelemetry.api.trace.Tracer;
+import io.opentelemetry.context.Context;
+
+public class FetchContext extends TraceContext {
+ public static final FetchContext DEFAULT = new FetchContext();
+ private ReadOptions readOptions = ReadOptions.DEFAULT;
+
+ public FetchContext() {
+ super(false, null, null);
+ }
+
+ public FetchContext(TraceContext context) {
+ super(context);
+ }
+
+ public FetchContext(boolean isTraceEnabled, Tracer tracer, Context currentContext) {
+ super(isTraceEnabled, tracer, currentContext);
+ }
+
+ public ReadOptions readOptions() {
+ return readOptions;
+ }
+
+ public void setReadOptions(ReadOptions readOptions) {
+ this.readOptions = readOptions;
+ }
+}
diff --git a/s3stream/src/main/java/com/automq/stream/s3/trace/AttributeBindings.java b/s3stream/src/main/java/com/automq/stream/s3/trace/AttributeBindings.java
new file mode 100644
index 000000000..c60247aeb
--- /dev/null
+++ b/s3stream/src/main/java/com/automq/stream/s3/trace/AttributeBindings.java
@@ -0,0 +1,106 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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
+ *
+ * http://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 com.automq.stream.s3.trace;
+
+import io.opentelemetry.api.common.AttributeKey;
+import io.opentelemetry.api.common.AttributesBuilder;
+import io.opentelemetry.instrumentation.annotations.SpanAttribute;
+
+import java.lang.reflect.Method;
+import java.lang.reflect.Parameter;
+import java.lang.reflect.Type;
+import java.util.function.BiFunction;
+
+public class AttributeBindings {
+ private final BiFunction[] bindings;
+
+ private AttributeBindings(BiFunction[] bindings) {
+ this.bindings = bindings;
+ }
+
+ public static AttributeBindings bind(Method method, String[] parametersNames) {
+ Parameter[] parameters = method.getParameters();
+ if (parameters.length != parametersNames.length) {
+ return new AttributeBindings(null);
+ }
+
+ BiFunction[] bindings = new BiFunction[parametersNames.length];
+ for (int i = 0; i < parametersNames.length; i++) {
+ Parameter parameter = parameters[i];
+
+ SpanAttribute spanAttribute = parameter.getAnnotation(SpanAttribute.class);
+ if (spanAttribute == null) {
+ bindings[i] = emptyBinding();
+ } else {
+ String attributeName = spanAttribute.value().isEmpty() ? parametersNames[i] : spanAttribute.value();
+ bindings[i] = createBinding(attributeName, parameter.getParameterizedType());
+ }
+ }
+ return new AttributeBindings(bindings);
+ }
+
+ public boolean isEmpty() {
+ return bindings == null || bindings.length == 0;
+ }
+
+ public void apply(AttributesBuilder target, Object[] args) {
+ if (args.length != bindings.length) {
+ return;
+ }
+
+ for (int i = 0; i < args.length; i++) {
+ bindings[i].apply(target, args[i]);
+ }
+ }
+
+ static BiFunction emptyBinding() {
+ return (builder, arg) -> builder;
+ }
+
+ static BiFunction createBinding(String name, Type type) {
+ // Simple scalar parameter types
+ if (type == String.class) {
+ AttributeKey key = AttributeKey.stringKey(name);
+ return (builder, arg) -> builder.put(key, (String) arg);
+ }
+ if (type == long.class || type == Long.class) {
+ AttributeKey key = AttributeKey.longKey(name);
+ return (builder, arg) -> builder.put(key, (Long) arg);
+ }
+ if (type == double.class || type == Double.class) {
+ AttributeKey key = AttributeKey.doubleKey(name);
+ return (builder, arg) -> builder.put(key, (Double) arg);
+ }
+ if (type == boolean.class || type == Boolean.class) {
+ AttributeKey key = AttributeKey.booleanKey(name);
+ return (builder, arg) -> builder.put(key, (Boolean) arg);
+ }
+ if (type == int.class || type == Integer.class) {
+ AttributeKey key = AttributeKey.longKey(name);
+ return (builder, arg) -> builder.put(key, ((Integer) arg).longValue());
+ }
+ if (type == float.class || type == Float.class) {
+ AttributeKey key = AttributeKey.doubleKey(name);
+ return (builder, arg) -> builder.put(key, ((Float) arg).doubleValue());
+ }
+
+ // Default parameter types
+ AttributeKey key = AttributeKey.stringKey(name);
+ return (builder, arg) -> builder.put(key, arg.toString());
+ }
+}
diff --git a/s3stream/src/main/java/com/automq/stream/s3/trace/MethodCache.java b/s3stream/src/main/java/com/automq/stream/s3/trace/MethodCache.java
new file mode 100644
index 000000000..c0151593c
--- /dev/null
+++ b/s3stream/src/main/java/com/automq/stream/s3/trace/MethodCache.java
@@ -0,0 +1,34 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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
+ *
+ * http://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 com.automq.stream.s3.trace;
+
+import java.lang.reflect.Method;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.function.Function;
+
+final class MethodCache extends ClassValue> {
+ public V computeIfAbsent(Method key, Function mappingFunction) {
+ return this.get(key.getDeclaringClass()).computeIfAbsent(key, mappingFunction);
+ }
+
+ @Override
+ protected Map computeValue(Class type) {
+ return new ConcurrentHashMap<>();
+ }
+}
diff --git a/s3stream/src/main/java/com/automq/stream/s3/trace/SpanAttributesExtractor.java b/s3stream/src/main/java/com/automq/stream/s3/trace/SpanAttributesExtractor.java
new file mode 100644
index 000000000..ffa4c0c2d
--- /dev/null
+++ b/s3stream/src/main/java/com/automq/stream/s3/trace/SpanAttributesExtractor.java
@@ -0,0 +1,46 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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
+ *
+ * http://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 com.automq.stream.s3.trace;
+
+import io.opentelemetry.api.common.Attributes;
+import io.opentelemetry.api.common.AttributesBuilder;
+
+import java.lang.reflect.Method;
+
+public final class SpanAttributesExtractor {
+
+ private final MethodCache cache;
+
+ public static SpanAttributesExtractor create() {
+ return new SpanAttributesExtractor(new MethodCache<>());
+ }
+
+ SpanAttributesExtractor(MethodCache cache) {
+ this.cache = cache;
+ }
+
+ public Attributes extract(Method method, String[] parametersNames, Object[] args) {
+ AttributesBuilder attributes = Attributes.builder();
+ AttributeBindings bindings =
+ cache.computeIfAbsent(method, (Method m) -> AttributeBindings.bind(m, parametersNames));
+ if (!bindings.isEmpty()) {
+ bindings.apply(attributes, args);
+ }
+ return attributes.build();
+ }
+}
\ No newline at end of file
diff --git a/s3stream/src/main/java/com/automq/stream/s3/trace/TraceUtils.java b/s3stream/src/main/java/com/automq/stream/s3/trace/TraceUtils.java
new file mode 100644
index 000000000..4f94724d1
--- /dev/null
+++ b/s3stream/src/main/java/com/automq/stream/s3/trace/TraceUtils.java
@@ -0,0 +1,143 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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
+ *
+ * http://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 com.automq.stream.s3.trace;
+
+import com.automq.stream.s3.trace.context.TraceContext;
+import io.opentelemetry.api.common.Attributes;
+import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.api.trace.StatusCode;
+import io.opentelemetry.api.trace.Tracer;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.instrumentation.annotations.WithSpan;
+import org.aspectj.lang.ProceedingJoinPoint;
+import org.aspectj.lang.reflect.MethodSignature;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.lang.reflect.Method;
+import java.util.concurrent.Callable;
+import java.util.concurrent.CompletableFuture;
+
+public class TraceUtils {
+ private static final Logger LOGGER = LoggerFactory.getLogger(TraceUtils.class);
+ private static final SpanAttributesExtractor EXTRACTOR = SpanAttributesExtractor.create();
+
+ public static Object trace(TraceContext context, ProceedingJoinPoint joinPoint, WithSpan withSpan) throws Throwable {
+ if (context.isTraceDisabled()) {
+ return joinPoint.proceed();
+ }
+
+ MethodSignature signature = (MethodSignature) joinPoint.getSignature();
+ Method method = signature.getMethod();
+ Object[] args = joinPoint.getArgs();
+
+ String className = method.getDeclaringClass().getSimpleName();
+ String spanName = withSpan.value().isEmpty() ? className + "::" + method.getName() : withSpan.value();
+
+ TraceContext.Scope scope = createAndStartSpan(context, spanName);
+ if (scope == null) {
+ return joinPoint.proceed();
+ }
+ Span span = scope.getSpan();
+ Attributes attributes = EXTRACTOR.extract(method, signature.getParameterNames(), args);
+ span.setAllAttributes(attributes);
+
+ try {
+ if (method.getReturnType() == CompletableFuture.class) {
+ return doTraceWhenReturnCompletableFuture(scope, joinPoint);
+ } else {
+ return doTraceWhenReturnObject(scope, joinPoint);
+ }
+ } catch (Throwable t) {
+ endSpan(scope, t);
+ throw t;
+ }
+ }
+
+ public static T runWithSpanSync(TraceContext context, Attributes attributes, String spanName, Callable callable) throws Throwable {
+ TraceContext.Scope scope = createAndStartSpan(context, spanName);
+ if (scope == null) {
+ return callable.call();
+ }
+ scope.getSpan().setAllAttributes(attributes);
+ try (scope) {
+ T ret = callable.call();
+ endSpan(scope, null);
+ return ret;
+ } catch (Throwable t) {
+ endSpan(scope, t);
+ throw t;
+ }
+ }
+
+ public static CompletableFuture runWithSpanAsync(TraceContext context, Attributes attributes, String spanName,
+ Callable> callable) throws Throwable {
+ TraceContext.Scope scope = createAndStartSpan(context, spanName);
+ if (scope == null) {
+ return callable.call();
+ }
+ scope.getSpan().setAllAttributes(attributes);
+ try (scope) {
+ CompletableFuture cf = callable.call();
+ cf.whenComplete((nil, ex) -> endSpan(scope, ex));
+ return cf;
+ } catch (Throwable t) {
+ endSpan(scope, t);
+ throw t;
+ }
+ }
+
+ public static TraceContext.Scope createAndStartSpan(TraceContext context, String name) {
+ if (context.isTraceDisabled()) {
+ return null;
+ }
+ Tracer tracer = context.tracer();
+ Context parentContext = context.currentContext();
+ Span span = tracer.spanBuilder(name)
+ .setParent(parentContext)
+ .startSpan();
+
+ return context.attachContext(parentContext.with(span));
+ }
+
+ public static void endSpan(TraceContext.Scope scope, Throwable t) {
+ if (scope == null) {
+ return;
+ }
+ if (t != null) {
+ scope.getSpan().recordException(t);
+ scope.getSpan().setStatus(StatusCode.ERROR, t.getMessage());
+ } else {
+ scope.getSpan().setStatus(StatusCode.OK);
+ }
+ scope.getSpan().end();
+ scope.close();
+ }
+
+ private static CompletableFuture doTraceWhenReturnCompletableFuture(TraceContext.Scope scope, ProceedingJoinPoint joinPoint) throws Throwable {
+ CompletableFuture future = (CompletableFuture) joinPoint.proceed();
+ return future.whenComplete((r, t) -> endSpan(scope, t));
+ }
+
+ private static Object doTraceWhenReturnObject(TraceContext.Scope scope, ProceedingJoinPoint joinPoint) throws Throwable {
+ Object result = joinPoint.proceed();
+ endSpan(scope, null);
+ return result;
+ }
+
+}
diff --git a/s3stream/src/main/java/com/automq/stream/s3/trace/aop/S3StreamTraceAspect.java b/s3stream/src/main/java/com/automq/stream/s3/trace/aop/S3StreamTraceAspect.java
new file mode 100644
index 000000000..3490797fb
--- /dev/null
+++ b/s3stream/src/main/java/com/automq/stream/s3/trace/aop/S3StreamTraceAspect.java
@@ -0,0 +1,44 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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
+ *
+ * http://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 com.automq.stream.s3.trace.aop;
+
+import com.automq.stream.s3.trace.TraceUtils;
+import com.automq.stream.s3.trace.context.TraceContext;
+import io.opentelemetry.instrumentation.annotations.WithSpan;
+import org.aspectj.lang.ProceedingJoinPoint;
+import org.aspectj.lang.annotation.Around;
+import org.aspectj.lang.annotation.Aspect;
+import org.aspectj.lang.annotation.Pointcut;
+
+@Aspect
+public class S3StreamTraceAspect {
+
+ @Pointcut("@annotation(withSpan)")
+ public void trace(WithSpan withSpan) {
+ }
+
+ @Around(value = "trace(withSpan) && execution(* com.automq.stream..*(..))", argNames = "joinPoint,withSpan")
+ public Object createSpan(ProceedingJoinPoint joinPoint, WithSpan withSpan) throws Throwable {
+ Object[] args = joinPoint.getArgs();
+ if (args.length > 0 && args[0] instanceof TraceContext context) {
+ return TraceUtils.trace(context, joinPoint, withSpan);
+ }
+
+ return joinPoint.proceed();
+ }
+}
diff --git a/s3stream/src/main/java/com/automq/stream/s3/trace/context/TraceContext.java b/s3stream/src/main/java/com/automq/stream/s3/trace/context/TraceContext.java
new file mode 100644
index 000000000..a8342221a
--- /dev/null
+++ b/s3stream/src/main/java/com/automq/stream/s3/trace/context/TraceContext.java
@@ -0,0 +1,95 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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
+ *
+ * http://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 com.automq.stream.s3.trace.context;
+
+import io.opentelemetry.api.GlobalOpenTelemetry;
+import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.api.trace.Tracer;
+import io.opentelemetry.context.Context;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import javax.annotation.concurrent.NotThreadSafe;
+
+/**
+ * Trace context that holds the current trace context. This class is not thread safe and should be copied before
+ * asynchronous usage.
+ */
+@NotThreadSafe
+public class TraceContext {
+ private static final Logger LOGGER = LoggerFactory.getLogger(TraceContext.class);
+ public static final TraceContext DEFAULT = new TraceContext(false, null, null);
+ private final boolean isTraceEnabled;
+ private final Tracer tracer;
+ private Context currContext;
+
+ public TraceContext(boolean isTraceEnabled, Tracer tracer, Context currContext) {
+ this.isTraceEnabled = isTraceEnabled;
+ if (isTraceEnabled && tracer == null) {
+ this.tracer = GlobalOpenTelemetry.getTracer("s3stream");
+ } else {
+ this.tracer = tracer;
+ }
+ if (isTraceEnabled && currContext == null) {
+ this.currContext = Context.current();
+ } else {
+ this.currContext = currContext;
+ }
+ }
+
+ public TraceContext(TraceContext traceContext) {
+ this(traceContext.isTraceEnabled, traceContext.tracer, traceContext.currContext);
+ }
+
+ public boolean isTraceDisabled() {
+ return !isTraceEnabled;
+ }
+
+ public Tracer tracer() {
+ return tracer;
+ }
+
+ public Context currentContext() {
+ return currContext;
+ }
+
+ public Scope attachContext(Context contextToAttach) {
+ return new Scope(contextToAttach);
+ }
+
+ public class Scope implements AutoCloseable {
+ private final Context prevContext;
+ private final Span span;
+
+ private Scope(Context contextToAttach) {
+ this.prevContext = currContext;
+ this.span = Span.fromContext(contextToAttach);
+ currContext = contextToAttach;
+ }
+
+ public Span getSpan() {
+ return span;
+ }
+
+ @Override
+ public void close() {
+ currContext = prevContext;
+ }
+ }
+
+}
diff --git a/s3stream/src/main/java/com/automq/stream/s3/wal/BlockWALService.java b/s3stream/src/main/java/com/automq/stream/s3/wal/BlockWALService.java
index 6fe679474..6e2ccfbdb 100644
--- a/s3stream/src/main/java/com/automq/stream/s3/wal/BlockWALService.java
+++ b/s3stream/src/main/java/com/automq/stream/s3/wal/BlockWALService.java
@@ -23,6 +23,8 @@
import com.automq.stream.s3.metrics.TimerUtil;
import com.automq.stream.s3.metrics.operations.S3Operation;
import com.automq.stream.s3.metrics.operations.S3Stage;
+import com.automq.stream.s3.trace.TraceUtils;
+import com.automq.stream.s3.trace.context.TraceContext;
import com.automq.stream.s3.wal.util.WALChannel;
import com.automq.stream.s3.wal.util.WALUtil;
import com.automq.stream.utils.ThreadUtils;
@@ -380,13 +382,18 @@ public WALMetadata metadata() {
}
@Override
- public AppendResult append(ByteBuf buf, int crc) throws OverCapacityException {
+ public AppendResult append(TraceContext context, ByteBuf buf, int crc) throws OverCapacityException {
+ // get current method name
+ TraceContext.Scope scope = TraceUtils.createAndStartSpan(context, "BlockWALService::append");
TimerUtil timerUtil = new TimerUtil();
try {
- return append0(buf, crc);
+ AppendResult result = append0(buf, crc);
+ result.future().whenComplete((nil, ex) -> TraceUtils.endSpan(scope, ex));
+ return result;
} catch (OverCapacityException ex) {
buf.release();
S3StreamMetricsManager.recordOperationLatency(timerUtil.elapsedAs(TimeUnit.NANOSECONDS), S3Operation.APPEND_STORAGE_WAL_FULL);
+ TraceUtils.endSpan(scope, ex);
throw ex;
}
}
diff --git a/s3stream/src/main/java/com/automq/stream/s3/wal/MemoryWriteAheadLog.java b/s3stream/src/main/java/com/automq/stream/s3/wal/MemoryWriteAheadLog.java
index ffb8601b6..9245b0a1c 100644
--- a/s3stream/src/main/java/com/automq/stream/s3/wal/MemoryWriteAheadLog.java
+++ b/s3stream/src/main/java/com/automq/stream/s3/wal/MemoryWriteAheadLog.java
@@ -17,6 +17,7 @@
package com.automq.stream.s3.wal;
+import com.automq.stream.s3.trace.context.TraceContext;
import io.netty.buffer.ByteBuf;
import java.io.IOException;
@@ -45,7 +46,7 @@ public WALMetadata metadata() {
}
@Override
- public AppendResult append(ByteBuf data, int crc) {
+ public AppendResult append(TraceContext traceContext, ByteBuf data, int crc) {
data.release();
long offset = offsetAlloc.getAndIncrement();
return new AppendResult() {
diff --git a/s3stream/src/main/java/com/automq/stream/s3/wal/WriteAheadLog.java b/s3stream/src/main/java/com/automq/stream/s3/wal/WriteAheadLog.java
index 6297b614a..4b301464a 100644
--- a/s3stream/src/main/java/com/automq/stream/s3/wal/WriteAheadLog.java
+++ b/s3stream/src/main/java/com/automq/stream/s3/wal/WriteAheadLog.java
@@ -18,6 +18,7 @@
package com.automq.stream.s3.wal;
+import com.automq.stream.s3.trace.context.TraceContext;
import io.netty.buffer.ByteBuf;
import java.io.IOException;
@@ -45,10 +46,18 @@ public interface WriteAheadLog {
*
* @return The data position will be written.
*/
- AppendResult append(ByteBuf data, int crc) throws OverCapacityException;
+ AppendResult append(TraceContext context, ByteBuf data, int crc) throws OverCapacityException;
+
+ default AppendResult append(TraceContext context, ByteBuf data) throws OverCapacityException {
+ return append(context, data, 0);
+ }
+
+ default AppendResult append(ByteBuf data, int crc) throws OverCapacityException {
+ return append(TraceContext.DEFAULT, data, crc);
+ }
default AppendResult append(ByteBuf data) throws OverCapacityException {
- return append(data, 0);
+ return append(TraceContext.DEFAULT, data, 0);
}
Iterator recover();
diff --git a/s3stream/src/test/java/com/automq/stream/s3/S3StorageTest.java b/s3stream/src/test/java/com/automq/stream/s3/S3StorageTest.java
index d36ae9f52..878b3774b 100644
--- a/s3stream/src/test/java/com/automq/stream/s3/S3StorageTest.java
+++ b/s3stream/src/test/java/com/automq/stream/s3/S3StorageTest.java
@@ -17,7 +17,6 @@
package com.automq.stream.s3;
-import com.automq.stream.api.ReadOptions;
import com.automq.stream.s3.cache.DefaultS3BlockCache;
import com.automq.stream.s3.cache.LogCache;
import com.automq.stream.s3.cache.ReadDataBlock;
@@ -103,9 +102,9 @@ public void testAppend() throws Exception {
cf2.get(3, TimeUnit.SECONDS);
cf3.get(3, TimeUnit.SECONDS);
- ReadDataBlock readRst = storage.read(233, 10, 13, 90, ReadOptions.DEFAULT).get();
+ ReadDataBlock readRst = storage.read(233, 10, 13, 90).get();
assertEquals(1, readRst.getRecords().size());
- readRst = storage.read(233, 10, 13, 200, ReadOptions.DEFAULT).get();
+ readRst = storage.read(233, 10, 13, 200).get();
assertEquals(2, readRst.getRecords().size());
storage.forceUpload(233L).get();
@@ -252,7 +251,7 @@ public void testRecoverContinuousRecords() {
public void testWALOverCapacity() throws WriteAheadLog.OverCapacityException {
storage.append(newRecord(233L, 10L));
storage.append(newRecord(233L, 11L));
- doThrow(new WriteAheadLog.OverCapacityException("test")).when(wal).append(any());
+ doThrow(new WriteAheadLog.OverCapacityException("test")).when(wal).append(any(), any());
Mockito.when(objectManager.prepareObject(eq(1), anyLong())).thenReturn(CompletableFuture.completedFuture(16L));
CommitStreamSetObjectResponse resp = new CommitStreamSetObjectResponse();
diff --git a/s3stream/src/test/java/com/automq/stream/s3/S3StreamTest.java b/s3stream/src/test/java/com/automq/stream/s3/S3StreamTest.java
index 21c83c386..3779731be 100644
--- a/s3stream/src/test/java/com/automq/stream/s3/S3StreamTest.java
+++ b/s3stream/src/test/java/com/automq/stream/s3/S3StreamTest.java
@@ -56,7 +56,7 @@ public void setup() {
@Test
public void testFetch() throws Throwable {
stream.confirmOffset.set(120L);
- Mockito.when(storage.read(eq(233L), eq(110L), eq(120L), eq(100), any()))
+ Mockito.when(storage.read(any(), eq(233L), eq(110L), eq(120L), eq(100)))
.thenReturn(CompletableFuture.completedFuture(newReadDataBlock(110, 115, 110)));
FetchResult rst = stream.fetch(110, 120, 100).get(1, TimeUnit.SECONDS);
assertEquals(1, rst.recordBatchList().size());
diff --git a/s3stream/src/test/java/com/automq/stream/s3/cache/StreamReaderTest.java b/s3stream/src/test/java/com/automq/stream/s3/cache/StreamReaderTest.java
index 94b5003fe..01a5a37e1 100644
--- a/s3stream/src/test/java/com/automq/stream/s3/cache/StreamReaderTest.java
+++ b/s3stream/src/test/java/com/automq/stream/s3/cache/StreamReaderTest.java
@@ -29,6 +29,7 @@
import com.automq.stream.s3.objects.ObjectManager;
import com.automq.stream.s3.operator.MemoryS3Operator;
import com.automq.stream.s3.operator.S3Operator;
+import com.automq.stream.s3.trace.context.TraceContext;
import com.automq.stream.utils.CloseableIterator;
import com.automq.stream.utils.Threads;
import org.apache.commons.lang3.tuple.ImmutablePair;
@@ -77,7 +78,7 @@ public void testGetDataBlockIndices() {
StreamReader streamReader = new StreamReader(s3Operator, objectManager, Mockito.mock(BlockCache.class), new HashMap<>(), new InflightReadThrottle());
StreamReader.ReadContext context = new StreamReader.ReadContext(15L, 1024);
- streamReader.getDataBlockIndices(233L, 1024L, context).thenAccept(v -> {
+ streamReader.getDataBlockIndices(TraceContext.DEFAULT, 233L, 1024L, context).thenAccept(v -> {
Assertions.assertEquals(40L, context.nextStartOffset);
Assertions.assertEquals(0, context.nextMaxBytes);
Assertions.assertEquals(2, context.streamDataBlocksPair.size());
@@ -115,7 +116,7 @@ public void testSyncReadAheadInflight() {
List.of(new StreamDataBlock(streamId, 64, 128, objectId, index1))))).when(reader).find(eq(streamId), eq(startOffset), anyLong(), eq(maxBytes));
doReturn(new CompletableFuture<>()).when(reader).read(index1);
- streamReader.syncReadAhead(streamId, startOffset, endOffset, maxBytes, Mockito.mock(ReadAheadAgent.class), UUID.randomUUID());
+ streamReader.syncReadAhead(TraceContext.DEFAULT, streamId, startOffset, endOffset, maxBytes, Mockito.mock(ReadAheadAgent.class), UUID.randomUUID());
Threads.sleep(1000);
Assertions.assertEquals(2, inflightReadAheadTasks.size());
ReadAheadTaskKey key1 = new ReadAheadTaskKey(233L, startOffset);
@@ -171,7 +172,7 @@ public StreamRecordBatch next() {
});
Mockito.when(reader.read(index1)).thenReturn(CompletableFuture.completedFuture(dataBlock1));
context.objectReaderMap = new HashMap<>(Map.of(1L, reader));
- CompletableFuture> cf = streamReader.handleSyncReadAhead(233L, 0,
+ CompletableFuture> cf = streamReader.handleSyncReadAhead(TraceContext.DEFAULT,233L, 0,
999, 64, Mockito.mock(ReadAheadAgent.class), UUID.randomUUID(), new TimerUtil(), context);
cf.whenComplete((rst, ex) -> {
@@ -233,7 +234,7 @@ public StreamRecordBatch next() {
ReadAheadTaskKey key = new ReadAheadTaskKey(233L, startOffset);
context.taskKeySet.add(key);
inflightReadAheadTasks.put(key, new DefaultS3BlockCache.ReadAheadTaskContext(new CompletableFuture<>(), DefaultS3BlockCache.ReadBlockCacheStatus.INIT));
- CompletableFuture> cf = streamReader.handleSyncReadAhead(233L, startOffset,
+ CompletableFuture> cf = streamReader.handleSyncReadAhead(TraceContext.DEFAULT, 233L, startOffset,
999, 64, Mockito.mock(ReadAheadAgent.class), UUID.randomUUID(), new TimerUtil(), context);
cf.whenComplete((rst, ex) -> {
@@ -294,7 +295,7 @@ public StreamRecordBatch next() {
Mockito.when(reader.read(index1)).thenReturn(CompletableFuture.completedFuture(dataBlock1));
Mockito.when(reader.read(index2)).thenReturn(CompletableFuture.failedFuture(new RuntimeException("exception")));
context.objectReaderMap = new HashMap<>(Map.of(1L, reader));
- CompletableFuture> cf = streamReader.handleSyncReadAhead(233L, 0,
+ CompletableFuture> cf = streamReader.handleSyncReadAhead(TraceContext.DEFAULT, 233L, 0,
512, 1024, Mockito.mock(ReadAheadAgent.class), UUID.randomUUID(), new TimerUtil(), context);
Threads.sleep(1000);
diff --git a/store/src/main/java/com/automq/rocketmq/store/S3StreamStore.java b/store/src/main/java/com/automq/rocketmq/store/S3StreamStore.java
index ef89bdaa6..c88aa9150 100644
--- a/store/src/main/java/com/automq/rocketmq/store/S3StreamStore.java
+++ b/store/src/main/java/com/automq/rocketmq/store/S3StreamStore.java
@@ -22,6 +22,7 @@
import com.automq.rocketmq.metadata.api.StoreMetadataService;
import com.automq.rocketmq.store.api.StreamStore;
import com.automq.rocketmq.store.model.StoreContext;
+import com.automq.rocketmq.store.util.ContextUtil;
import com.automq.stream.api.AppendResult;
import com.automq.stream.api.FetchResult;
import com.automq.stream.api.OpenStreamOptions;
@@ -35,6 +36,8 @@
import com.automq.stream.s3.cache.DefaultS3BlockCache;
import com.automq.stream.s3.cache.S3BlockCache;
import com.automq.stream.s3.compact.CompactionManager;
+import com.automq.stream.s3.context.AppendContext;
+import com.automq.stream.s3.context.FetchContext;
import com.automq.stream.s3.network.AsyncNetworkBandwidthLimiter;
import com.automq.stream.s3.objects.ObjectManager;
import com.automq.stream.s3.operator.DefaultS3Operator;
@@ -127,7 +130,8 @@ public CompletableFuture fetch(StoreContext context, @SpanAttribute
if (stream.isEmpty()) {
throw new IllegalStateException("Stream " + streamId + " is not opened.");
}
- return stream.get().fetch(startOffset, startOffset + maxCount, Integer.MAX_VALUE)
+ FetchContext fetchContext = new FetchContext(ContextUtil.buildStreamTraceContext(context));
+ return stream.get().fetch(fetchContext, startOffset, startOffset + maxCount, Integer.MAX_VALUE)
.thenApplyAsync(result -> {
context.span().ifPresent(span -> {
span.setAttribute("messageCount", result.recordBatchList().size());
@@ -151,7 +155,8 @@ public CompletableFuture append(StoreContext context, long streamI
span.setAttribute("recordBytes", recordBatch.rawPayload().remaining());
});
- return stream.get().append(recordBatch)
+ AppendContext appendContext = new AppendContext(ContextUtil.buildStreamTraceContext(context));
+ return stream.get().append(appendContext, recordBatch)
.thenApplyAsync(result -> {
context.span().ifPresent(span -> span.setAttribute("offset", result.baseOffset()));
return result;
diff --git a/store/src/main/java/com/automq/rocketmq/store/util/ContextUtil.java b/store/src/main/java/com/automq/rocketmq/store/util/ContextUtil.java
new file mode 100644
index 000000000..b079ac38d
--- /dev/null
+++ b/store/src/main/java/com/automq/rocketmq/store/util/ContextUtil.java
@@ -0,0 +1,29 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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
+ *
+ * http://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 com.automq.rocketmq.store.util;
+
+import com.automq.stream.s3.trace.context.TraceContext;
+import io.opentelemetry.context.Context;
+
+public class ContextUtil {
+ public static TraceContext buildStreamTraceContext(com.automq.rocketmq.common.trace.TraceContext context) {
+ boolean isTraceEnabled = context.tracer().isPresent();
+ Context currContext = context.span().map(span -> Context.current().with(span)).orElse(Context.current());
+ return new TraceContext(isTraceEnabled, context.tracer().get(), currContext);
+ }
+}
diff --git a/store/src/test/java/com/automq/rocketmq/store/mock/MemoryStreamClient.java b/store/src/test/java/com/automq/rocketmq/store/mock/MemoryStreamClient.java
index ddf5e12e0..eab43557f 100644
--- a/store/src/test/java/com/automq/rocketmq/store/mock/MemoryStreamClient.java
+++ b/store/src/test/java/com/automq/rocketmq/store/mock/MemoryStreamClient.java
@@ -21,11 +21,13 @@
import com.automq.stream.api.CreateStreamOptions;
import com.automq.stream.api.FetchResult;
import com.automq.stream.api.OpenStreamOptions;
-import com.automq.stream.api.ReadOptions;
import com.automq.stream.api.RecordBatch;
import com.automq.stream.api.RecordBatchWithContext;
import com.automq.stream.api.Stream;
import com.automq.stream.api.StreamClient;
+import com.automq.stream.s3.context.AppendContext;
+import com.automq.stream.s3.context.FetchContext;
+
import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.List;
@@ -99,14 +101,14 @@ public long nextOffset() {
}
@Override
- public CompletableFuture append(RecordBatch recordBatch) {
+ public CompletableFuture append(AppendContext context, RecordBatch recordBatch) {
long baseOffset = nextOffsetAlloc.getAndAdd(recordBatch.count());
recordMap.put(baseOffset, new RecordBatchWithContextWrapper(recordBatch, baseOffset));
return CompletableFuture.completedFuture(() -> baseOffset);
}
@Override
- public CompletableFuture fetch(long startOffset, long endOffset, int maxBytesHint, ReadOptions readOptions) {
+ public CompletableFuture fetch(FetchContext context, long startOffset, long endOffset, int maxBytesHint) {
List records = new ArrayList<>(recordMap.subMap(startOffset, endOffset).values());
return CompletableFuture.completedFuture(() -> records);
}