diff --git a/core/src/main/java/kafka/automq/zerozone/DefaultRouterChannelProvider.java b/core/src/main/java/kafka/automq/zerozone/DefaultRouterChannelProvider.java index cff5f0e6f7..969d38f2a1 100644 --- a/core/src/main/java/kafka/automq/zerozone/DefaultRouterChannelProvider.java +++ b/core/src/main/java/kafka/automq/zerozone/DefaultRouterChannelProvider.java @@ -24,6 +24,7 @@ import org.apache.kafka.image.MetadataDelta; import org.apache.kafka.image.MetadataImage; +import com.automq.stream.s3.metrics.operations.S3RequestSource; import com.automq.stream.s3.network.GlobalNetworkBandwidthLimiters; import com.automq.stream.s3.operator.BucketURI; import com.automq.stream.s3.operator.ObjectStorage; @@ -103,6 +104,7 @@ private RouterChannel newChannel(int ownerNodeId, boolean readOnly) { .withBucketId(bucket.bucketId()) .withOpenMode(bucketReadOnly ? OpenMode.READ_ONLY : OpenMode.READ_WRITE) .withType(WAL_TYPE) + .withSource(S3RequestSource.ROUTER_WAL) .withManualMode(true); if (!bucketReadOnly) { builder.withEpoch(nodeEpoch); diff --git a/core/src/main/scala/kafka/log/stream/s3/wal/DefaultWalFactory.java b/core/src/main/scala/kafka/log/stream/s3/wal/DefaultWalFactory.java index 885abd8ffe..cbc2e2dcd9 100644 --- a/core/src/main/scala/kafka/log/stream/s3/wal/DefaultWalFactory.java +++ b/core/src/main/scala/kafka/log/stream/s3/wal/DefaultWalFactory.java @@ -20,6 +20,7 @@ package kafka.log.stream.s3.wal; import com.automq.shell.AutoMQApplication; +import com.automq.stream.s3.metrics.operations.S3RequestSource; import com.automq.stream.s3.network.NetworkBandwidthLimiter; import com.automq.stream.s3.operator.BucketURI; import com.automq.stream.s3.operator.ObjectStorage; @@ -56,6 +57,7 @@ public WriteAheadLog build(IdURI uri, BuildOptions options) { switch (uri.protocol().toUpperCase(Locale.ENGLISH)) { case "S3": ObjectWALConfig.Builder configBuilder = ObjectWALConfig.builder().withURI(uri) + .withSource(S3RequestSource.DELTA_WAL) .withClusterId(AutoMQApplication.getClusterId()) .withNodeId(nodeId) .withEpoch(options.nodeEpoch()) @@ -68,7 +70,8 @@ public WriteAheadLog build(IdURI uri, BuildOptions options) { .outboundLimiter(networkOutboundLimiter) .build(); - ReservationService reservationService = new ObjectReservationService(AutoMQApplication.getClusterId(), walObjectStorage, walObjectStorage.bucketId()); + ReservationService reservationService = new ObjectReservationService(AutoMQApplication.getClusterId(), + walObjectStorage, walObjectStorage.bucketId(), S3RequestSource.DELTA_WAL); configBuilder.withReservationService(reservationService); // This factory creates a dedicated ObjectStorage, so the WAL owns its lifecycle. return new ObjectWALService(Time.SYSTEM, walObjectStorage, configBuilder.build(), true); diff --git a/core/src/test/java/kafka/automq/zerozone/DefaultRouterChannelProviderTest.java b/core/src/test/java/kafka/automq/zerozone/DefaultRouterChannelProviderTest.java index eb512bce24..ad7cc4a812 100644 --- a/core/src/test/java/kafka/automq/zerozone/DefaultRouterChannelProviderTest.java +++ b/core/src/test/java/kafka/automq/zerozone/DefaultRouterChannelProviderTest.java @@ -19,6 +19,7 @@ package kafka.automq.zerozone; +import com.automq.stream.s3.metrics.operations.S3RequestSource; import com.automq.stream.s3.operator.BucketURI; import com.automq.stream.s3.operator.ObjectStorage; import com.automq.stream.s3.wal.OpenMode; @@ -68,6 +69,7 @@ public void testReadWriteChannelUsesURIOptions() { assertEquals(123L, writer.epoch()); assertEquals("cluster", writer.clusterId()); assertEquals("rc", writer.type()); + assertEquals(S3RequestSource.ROUTER_WAL, writer.source()); assertEquals(OpenMode.READ_WRITE, writer.openMode()); assertEquals("test-secret", IdURI.parse(writer.uri()).extensionString("secretKey")); assertFalse(writer.toString().contains("test-secret")); @@ -75,6 +77,7 @@ public void testReadWriteChannelUsesURIOptions() { ObjectWALConfig reader = configs.get(1); assertEquals(OpenMode.READ_ONLY, reader.openMode()); + assertEquals(S3RequestSource.ROUTER_WAL, reader.source()); assertEquals(43, reader.nodeId()); assertEquals(writer.uri(), reader.uri()); assertEquals(1, reader.version()); @@ -118,6 +121,7 @@ public void testMultipleBuckets() { provider.channel(); provider.readOnlyChannel(43); assertEquals(4, wals.constructed().size()); + configs.forEach(config -> assertEquals(S3RequestSource.ROUTER_WAL, config.source())); assertEquals(100, configs.get(0).batchInterval()); assertEquals(200, configs.get(1).batchInterval()); assertEquals(0, configs.get(0).bucketId()); diff --git a/core/src/test/java/kafka/log/stream/s3/wal/DefaultWalFactoryTest.java b/core/src/test/java/kafka/log/stream/s3/wal/DefaultWalFactoryTest.java index c0732f394c..774c1f43ef 100644 --- a/core/src/test/java/kafka/log/stream/s3/wal/DefaultWalFactoryTest.java +++ b/core/src/test/java/kafka/log/stream/s3/wal/DefaultWalFactoryTest.java @@ -19,11 +19,13 @@ package kafka.log.stream.s3.wal; +import com.automq.stream.s3.metrics.operations.S3RequestSource; import com.automq.stream.s3.network.NetworkBandwidthLimiter; import com.automq.stream.s3.operator.ObjectStorage; import com.automq.stream.s3.operator.ObjectStorageFactory; import com.automq.stream.s3.wal.OpenMode; import com.automq.stream.s3.wal.WalFactory; +import com.automq.stream.s3.wal.impl.object.ObjectReservationService; import com.automq.stream.s3.wal.impl.object.ObjectWALConfig; import com.automq.stream.s3.wal.impl.object.ObjectWALService; import com.automq.stream.utils.IdURI; @@ -58,6 +60,8 @@ public void testHistoricalURIRespectsLayoutVersion(OpenMode mode) { List configs = new ArrayList<>(); try (MockedConstruction builders = mockConstruction(ObjectStorageFactory.Builder.class, withSettings().defaultAnswer(RETURNS_SELF), (builder, context) -> when(builder.build()).thenReturn(mock(ObjectStorage.class))); + MockedConstruction reservations = mockConstruction(ObjectReservationService.class, + (reservation, context) -> assertEquals(S3RequestSource.DELTA_WAL, context.arguments().get(3))); MockedConstruction wals = mockConstruction(ObjectWALService.class, (wal, context) -> configs.add((ObjectWALConfig) context.arguments().get(2)))) { WalFactory.BuildOptions options = WalFactory.BuildOptions.builder().nodeEpoch(42).openMode(mode).build(); @@ -69,7 +73,9 @@ public void testHistoricalURIRespectsLayoutVersion(OpenMode mode) { assertEquals(0, configs.get(2).version()); assertEquals(3, builders.constructed().size()); assertEquals(3, wals.constructed().size()); + assertEquals(3, reservations.constructed().size()); configs.forEach(config -> assertEquals(mode, config.openMode())); + configs.forEach(config -> assertEquals(S3RequestSource.DELTA_WAL, config.source())); } } diff --git a/s3stream/src/main/java/com/automq/stream/s3/metrics/operations/S3RequestSource.java b/s3stream/src/main/java/com/automq/stream/s3/metrics/operations/S3RequestSource.java new file mode 100644 index 0000000000..57e3243a57 --- /dev/null +++ b/s3stream/src/main/java/com/automq/stream/s3/metrics/operations/S3RequestSource.java @@ -0,0 +1,33 @@ +/* + * Copyright 2026, AutoMQ HK Limited. + * + * 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.metrics.operations; + +/** + * Bounded ownership labels for S3 GET/PUT observations. Values are stable metric labels, + * not per-producer context; shared WAL helpers can also perform replay or maintenance I/O. + */ +public enum S3RequestSource { + /** GET/PUT callers that have not supplied source attribution. */ + UNCLASSIFIED, + ROUTER_WAL, + DELTA_WAL, + /** One merged GET serving child reads with different source labels. */ + MIXED +} diff --git a/s3stream/src/main/java/com/automq/stream/s3/metrics/stats/OperationLatencyMetrics.java b/s3stream/src/main/java/com/automq/stream/s3/metrics/stats/OperationLatencyMetrics.java index 8a66c4f559..f32105211b 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/metrics/stats/OperationLatencyMetrics.java +++ b/s3stream/src/main/java/com/automq/stream/s3/metrics/stats/OperationLatencyMetrics.java @@ -22,6 +22,7 @@ import com.automq.stream.s3.metrics.Metrics; import com.automq.stream.s3.metrics.MetricsLevel; import com.automq.stream.s3.metrics.operations.S3Operation; +import com.automq.stream.s3.metrics.operations.S3RequestSource; import com.automq.stream.s3.metrics.operations.S3Stage; import com.automq.stream.s3.metrics.wrapper.DeltaHistogram; @@ -66,6 +67,15 @@ public static DeltaHistogram operation(MetricsLevel metricsLevel, String operati operationAttributes(operationType, operationName, status, sizeLabelName)); } + /** Registers a sized S3 request histogram with one bounded initiating source. */ + public static DeltaHistogram request(MetricsLevel metricsLevel, S3Operation operation, String status, + String sizeLabelName, S3RequestSource source) { + return OPERATION_LATENCY.histogram(metricsLevel, Attributes.builder() + .putAll(operationAttributes(operation.getType().getName(), operation.getName(), status, sizeLabelName)) + .put("source", source.name()) + .build()); + } + public static DeltaHistogram stage(MetricsLevel metricsLevel, S3Stage stage) { return OPERATION_LATENCY.histogram(metricsLevel, Attributes.builder() .putAll(operationAttributes(stage.getOperation().getType().getName(), stage.getOperation().getName())) diff --git a/s3stream/src/main/java/com/automq/stream/s3/operator/AbstractObjectStorage.java b/s3stream/src/main/java/com/automq/stream/s3/operator/AbstractObjectStorage.java index 64b9278ec5..7a66dd35d6 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/operator/AbstractObjectStorage.java +++ b/s3stream/src/main/java/com/automq/stream/s3/operator/AbstractObjectStorage.java @@ -22,6 +22,7 @@ import com.automq.stream.s3.exceptions.ObjectNotExistException; import com.automq.stream.s3.metrics.TimerUtil; import com.automq.stream.s3.metrics.operations.S3Operation; +import com.automq.stream.s3.metrics.operations.S3RequestSource; import com.automq.stream.s3.network.NetworkBandwidthLimiter; import com.automq.stream.s3.network.ThrottleStrategy; import com.automq.stream.s3.objects.ObjectAttributes; @@ -203,8 +204,9 @@ public Writer writer(WriteOptions options, String objectPath) { @Override public CompletableFuture rangeRead(ReadOptions options, String objectPath, long start, long end) { + ReadOptions requestOptions = options.copy(); CompletableFuture cf = new CompletableFuture<>(); - if (!bucketCheck(options.bucket(), cf)) { + if (!bucketCheck(requestOptions.bucket(), cf)) { return cf; } if (end != RANGE_READ_TO_END && start > end) { @@ -234,12 +236,12 @@ public CompletableFuture rangeRead(ReadOptions options, String objectPa }); } - networkInboundBandwidthLimiterFunction.apply(options.throttleStrategy(), acquiredSize).whenComplete((v, ex) -> { + networkInboundBandwidthLimiterFunction.apply(requestOptions.throttleStrategy(), acquiredSize).whenComplete((v, ex) -> { if (ex != null) { cf.completeExceptionally(ex); } else { synchronized (waitingReadTasks) { - waitingReadTasks.add(new AbstractObjectStorage.ReadTask(options, objectPath, start, end, cf)); + waitingReadTasks.add(new AbstractObjectStorage.ReadTask(requestOptions, objectPath, start, end, cf)); } } }); @@ -249,6 +251,7 @@ public CompletableFuture rangeRead(ReadOptions options, String objectPa @Override public CompletableFuture write(WriteOptions options, String objectPath, ByteBuf data) { + WriteOptions requestOptions = options.copy(); CompletableFuture cf = new CompletableFuture<>(); CompletableFuture retCf = acquireWritePermit(cf).thenApply(nil -> new WriteResult(bucketURI.bucketId())); if (retCf.isDone()) { @@ -256,18 +259,18 @@ public CompletableFuture write(WriteOptions options, String objectP return retCf; } networkOutboundBandwidthLimiter - .consume(options.throttleStrategy(), data.readableBytes()) + .consume(requestOptions.throttleStrategy(), data.readableBytes()) .whenCompleteAsync((v, ex) -> { if (ex != null) { cf.completeExceptionally(ex); data.release(); return; } - if (checkTimeout(options, cf)) { + if (checkTimeout(requestOptions, cf)) { data.release(); return; } - queuedWrite0(options, objectPath, data, cf); + queuedWrite0(requestOptions, objectPath, data, cf); }, writeLimiterCallbackExecutor); return retCf; } @@ -309,10 +312,10 @@ private void copy0(String sourceBucket, String sourcePath, String destinationPat }); } - private void recordWriteStats(String path, long objectSize, TimerUtil timerUtil) { + private void recordWriteStats(WriteOptions options, String path, long objectSize, TimerUtil timerUtil) { s3LatencyCalculator.record(objectSize, timerUtil.elapsedAs(TimeUnit.MILLISECONDS)); ObjectStorageMetrics.recordUploadSize(objectSize); - ObjectStorageMetrics.recordPutObject(objectSize, true, timerUtil.elapsedAs(TimeUnit.NANOSECONDS)); + ObjectStorageMetrics.recordPutObject(options.source(), objectSize, true, timerUtil.elapsedAs(TimeUnit.NANOSECONDS)); successWriteMonitor.record(objectSize); if (logger.isDebugEnabled()) { logger.debug("put object {} with size {}, cost {}ms", path, objectSize, timerUtil.elapsedAs(TimeUnit.NANOSECONDS)); @@ -355,7 +358,7 @@ private void write0(WriteOptions options, String path, ByteBuf data, Completable } writeCf.thenAccept(nil -> { - recordWriteStats(path, objectSize, timerUtil); + recordWriteStats(options, path, objectSize, timerUtil); data.release(); finalCf.complete(null); }).exceptionally(ex -> { @@ -368,7 +371,7 @@ private void write0(WriteOptions options, String path, ByteBuf data, Completable return null; } - ObjectStorageMetrics.recordPutObject(objectSize, false, timerUtil.elapsedAs(TimeUnit.NANOSECONDS)); + ObjectStorageMetrics.recordPutObject(options.source(), objectSize, false, timerUtil.elapsedAs(TimeUnit.NANOSECONDS)); Pair strategyAndCause = toRetryStrategyAndCause(ex, S3Operation.PUT_OBJECT); RetryStrategy retryStrategy = strategyAndCause.getLeft(); Throwable cause = strategyAndCause.getRight(); @@ -407,9 +410,9 @@ private CompletableFuture doFastRetryWrite(WriteOptions options, String pa } return cf.whenComplete((nil, ex) -> { if (ex == null) { - recordWriteStats(path, objectSize, timerUtil); + recordWriteStats(options, path, objectSize, timerUtil); } else if (!finalCf.isDone()) { - ObjectStorageMetrics.recordPutObject(objectSize, false, timerUtil.elapsedAs(TimeUnit.NANOSECONDS)); + ObjectStorageMetrics.recordPutObject(options.source(), objectSize, false, timerUtil.elapsedAs(TimeUnit.NANOSECONDS)); } }); } @@ -822,7 +825,7 @@ private void tryMergeRead0() { path, mergedReadTask.start, mergedReadTask.end, mergedReadTask.end - mergedReadTask.start, mergedReadTask.dataSparsityRate); } - mergedRangeRead(mergedReadTask.readTasks.get(0).options, path, mergedReadTask.start, mergedReadTask.end) + mergedRangeRead(mergedReadTask.options(), path, mergedReadTask.start, mergedReadTask.end) .whenComplete((rst, ex) -> FutureUtil.suppress(() -> mergedReadTask.handleReadCompleted(rst, ex), logger)); } ); @@ -853,7 +856,7 @@ CompletableFuture mergedRangeRead(ReadOptions options, String path, lon return retCf; } - private void recordReadStats(String path, long start, long end, ByteBuf buf, TimerUtil timerUtil) { + private void recordReadStats(ReadOptions options, String path, long start, long end, ByteBuf buf, TimerUtil timerUtil) { long dataSize = buf.readableBytes(); readLatencyCalculator.record(dataSize, timerUtil.elapsedAs(TimeUnit.MILLISECONDS)); if (logger.isDebugEnabled()) { @@ -861,7 +864,7 @@ private void recordReadStats(String path, long start, long end, ByteBuf buf, Tim path, start, end, dataSize, timerUtil.elapsedAs(TimeUnit.MILLISECONDS)); } ObjectStorageMetrics.recordDownloadSize(dataSize); - ObjectStorageMetrics.recordGetObject(dataSize, true, timerUtil.elapsedAs(TimeUnit.NANOSECONDS)); + ObjectStorageMetrics.recordGetObject(options.source(), dataSize, true, timerUtil.elapsedAs(TimeUnit.NANOSECONDS)); } private CompletableFuture doFastRetryRead(ReadOptions options, String path, long start, long end, @@ -876,13 +879,13 @@ private CompletableFuture doFastRetryRead(ReadOptions options, String p return cf.whenComplete((buf, ex) -> { if (ex == null) { try { - recordReadStats(path, start, end, buf, timerUtil); + recordReadStats(options, path, start, end, buf, timerUtil); } catch (Throwable e) { buf.release(); throw e; } } else if (!finalCf.isDone()) { - ObjectStorageMetrics.recordGetObject(end - start, false, timerUtil.elapsedAs(TimeUnit.NANOSECONDS)); + ObjectStorageMetrics.recordGetObject(options.source(), end - start, false, timerUtil.elapsedAs(TimeUnit.NANOSECONDS)); } }); } @@ -898,7 +901,7 @@ private CompletableFuture mergedRangeRead0(ReadOptions options, String readCf.thenAccept(buf -> { boolean transferred = false; try { - recordReadStats(path, start, end, buf, timerUtil); + recordReadStats(options, path, start, end, buf, timerUtil); transferred = cf.complete(buf); } finally { if (!transferred) { @@ -922,7 +925,7 @@ private CompletableFuture mergedRangeRead0(ReadOptions options, String logger.warn("GetObject for object {} [{}, {}) fail, retry in {}ms", path, start, end, delay, cause); scheduler.schedule(() -> mergedRangeRead0(options, path, start, end, cf), delay, TimeUnit.MILLISECONDS); } - ObjectStorageMetrics.recordGetObject(size, false, timerUtil.elapsedAs(TimeUnit.NANOSECONDS)); + ObjectStorageMetrics.recordGetObject(options.source(), size, false, timerUtil.elapsedAs(TimeUnit.NANOSECONDS)); return null; }); return readCf; @@ -1102,6 +1105,18 @@ static class MergedReadTask { this.maxMergeReadSparsityRate = maxMergeReadSparsityRate; } + // Preserve the existing first-task settings, but do not assign heterogeneous work to its first caller. + ReadOptions options() { + ReadOptions options = readTasks.get(0).options.copy(); + for (ReadTask task : readTasks) { + if (task.options.source() != options.source()) { + options.source(S3RequestSource.MIXED); + break; + } + } + return options; + } + boolean tryMerge(AbstractObjectStorage.ReadTask readTask) { if (!canMerge(readTask)) { return false; diff --git a/s3stream/src/main/java/com/automq/stream/s3/operator/ObjectStorage.java b/s3stream/src/main/java/com/automq/stream/s3/operator/ObjectStorage.java index 4506ca0964..982badef39 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/operator/ObjectStorage.java +++ b/s3stream/src/main/java/com/automq/stream/s3/operator/ObjectStorage.java @@ -21,6 +21,7 @@ import com.automq.stream.s3.ByteBufAlloc; import com.automq.stream.s3.exceptions.ObjectNotExistException; +import com.automq.stream.s3.metrics.operations.S3RequestSource; import com.automq.stream.s3.network.ThrottleStrategy; import java.util.List; @@ -313,6 +314,7 @@ class WriteOptions { public static final short UNSET_BUCKET = (short) -2; private ThrottleStrategy throttleStrategy = ThrottleStrategy.BYPASS; + private S3RequestSource source = S3RequestSource.UNCLASSIFIED; private int allocType = ByteBufAlloc.DEFAULT; // timeout for one single network rpc private long apiCallAttemptTimeout = -1L; @@ -326,6 +328,17 @@ class WriteOptions { private long requestTime = System.nanoTime(); // write context end + /** Sets the bounded metric source; it does not change storage behavior or object metadata. */ + public WriteOptions source(S3RequestSource source) { + this.source = Objects.requireNonNull(source, "source"); + return this; + } + + /** Returns the source carried by this request and its retries. */ + public S3RequestSource source() { + return source; + } + public WriteOptions throttleStrategy(ThrottleStrategy throttleStrategy) { this.throttleStrategy = throttleStrategy; return this; @@ -415,6 +428,7 @@ public long timeout() { public WriteOptions copy() { WriteOptions copy = new WriteOptions(); copy.throttleStrategy = throttleStrategy; + copy.source = source; copy.allocType = allocType; copy.apiCallAttemptTimeout = apiCallAttemptTimeout; copy.bucketId = bucketId; @@ -431,9 +445,21 @@ class ReadOptions { public static final short UNSET_BUCKET = (short) -2; private ThrottleStrategy throttleStrategy = ThrottleStrategy.BYPASS; + private S3RequestSource source = S3RequestSource.UNCLASSIFIED; private short bucket = UNSET_BUCKET; private int retryCount; + /** Sets the bounded metric source; it does not change storage behavior or object metadata. */ + public ReadOptions source(S3RequestSource source) { + this.source = Objects.requireNonNull(source, "source"); + return this; + } + + /** Returns the source carried by this request and its retries. */ + public S3RequestSource source() { + return source; + } + public ReadOptions throttleStrategy(ThrottleStrategy throttleStrategy) { this.throttleStrategy = throttleStrategy; return this; @@ -466,6 +492,7 @@ public int retryCount() { public ReadOptions copy() { ReadOptions copy = new ReadOptions(); copy.throttleStrategy = throttleStrategy; + copy.source = source; copy.bucket = bucket; copy.retryCount = retryCount; return copy; diff --git a/s3stream/src/main/java/com/automq/stream/s3/operator/ObjectStorageMetrics.java b/s3stream/src/main/java/com/automq/stream/s3/operator/ObjectStorageMetrics.java index 5bbaac8c8b..9495394fcb 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/operator/ObjectStorageMetrics.java +++ b/s3stream/src/main/java/com/automq/stream/s3/operator/ObjectStorageMetrics.java @@ -24,6 +24,7 @@ import com.automq.stream.s3.metrics.MetricsLevel; import com.automq.stream.s3.metrics.S3StreamMetricsConstant; import com.automq.stream.s3.metrics.operations.S3Operation; +import com.automq.stream.s3.metrics.operations.S3RequestSource; import com.automq.stream.s3.metrics.stats.OperationLatencyMetrics; import com.automq.stream.s3.metrics.wrapper.DeltaHistogram; @@ -50,10 +51,10 @@ final class ObjectStorageMetrics { .histogram("kafka_stream_read_s3_limiter_time", "Time blocked on waiting for inflight read quota", "nanoseconds"); private static final Metrics.HistogramBundle WRITE_S3_LIMITER_TIME = Metrics.instance() .histogram("kafka_stream_write_s3_limiter_time", "Time blocked on waiting for inflight write quota", "nanoseconds"); - private static final Map GET_OBJECT_SUCCESS_STATS = new ConcurrentHashMap<>(); - private static final Map GET_OBJECT_FAILED_STATS = new ConcurrentHashMap<>(); - private static final Map PUT_OBJECT_SUCCESS_STATS = new ConcurrentHashMap<>(); - private static final Map PUT_OBJECT_FAILED_STATS = new ConcurrentHashMap<>(); + private static final Map> GET_OBJECT_SUCCESS_STATS = new ConcurrentHashMap<>(); + private static final Map> GET_OBJECT_FAILED_STATS = new ConcurrentHashMap<>(); + private static final Map> PUT_OBJECT_SUCCESS_STATS = new ConcurrentHashMap<>(); + private static final Map> PUT_OBJECT_FAILED_STATS = new ConcurrentHashMap<>(); private static final DeltaHistogram LIST_OBJECTS_SUCCESS_STATS = operationStats( S3Operation.LIST_OBJECTS, S3StreamMetricsConstant.LABEL_STATUS_SUCCESS); private static final DeltaHistogram LIST_OBJECTS_FAILED_STATS = operationStats( @@ -94,13 +95,13 @@ static void recordDownloadSize(long size) { DOWNLOAD_SIZE_TOTAL.add(size); } - static void recordGetObject(long size, boolean success, long timeNanos) { - objectSizeStats(GET_OBJECT_SUCCESS_STATS, GET_OBJECT_FAILED_STATS, S3Operation.GET_OBJECT, size, success) + static void recordGetObject(S3RequestSource source, long size, boolean success, long timeNanos) { + requestStats(GET_OBJECT_SUCCESS_STATS, GET_OBJECT_FAILED_STATS, S3Operation.GET_OBJECT, source, size, success) .record(timeNanos); } - static void recordPutObject(long size, boolean success, long timeNanos) { - objectSizeStats(PUT_OBJECT_SUCCESS_STATS, PUT_OBJECT_FAILED_STATS, S3Operation.PUT_OBJECT, size, success) + static void recordPutObject(S3RequestSource source, long size, boolean success, long timeNanos) { + requestStats(PUT_OBJECT_SUCCESS_STATS, PUT_OBJECT_FAILED_STATS, S3Operation.PUT_OBJECT, source, size, success) .record(timeNanos); } @@ -135,6 +136,17 @@ static void recordWriteS3Limiter(int index, long timeNanos) { k -> WRITE_S3_LIMITER_TIME.histogram(MetricsLevel.DEBUG, attributes(index))).record(timeNanos); } + private static DeltaHistogram requestStats(Map> successStats, + Map> failedStats, S3Operation operation, S3RequestSource source, + long size, boolean success) { + Map stats = (success ? successStats : failedStats) + .computeIfAbsent(source, ignored -> new ConcurrentHashMap<>()); + String sizeLabel = AttributesUtils.getObjectBucketLabel(size); + return stats.computeIfAbsent(sizeLabel, label -> OperationLatencyMetrics.request(MetricsLevel.INFO, operation, + success ? S3StreamMetricsConstant.LABEL_STATUS_SUCCESS : S3StreamMetricsConstant.LABEL_STATUS_FAILED, + label, source)); + } + private static DeltaHistogram objectSizeStats(Map successStats, Map failedStats, S3Operation operation, long size, boolean success) { String label = AttributesUtils.getObjectBucketLabel(size); diff --git a/s3stream/src/main/java/com/automq/stream/s3/wal/impl/object/DefaultReader.java b/s3stream/src/main/java/com/automq/stream/s3/wal/impl/object/DefaultReader.java index 4a07574e95..285fcb85f0 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/wal/impl/object/DefaultReader.java +++ b/s3stream/src/main/java/com/automq/stream/s3/wal/impl/object/DefaultReader.java @@ -20,6 +20,7 @@ package com.automq.stream.s3.wal.impl.object; import com.automq.stream.s3.exceptions.ObjectNotExistException; +import com.automq.stream.s3.metrics.operations.S3RequestSource; import com.automq.stream.s3.model.StreamRecordBatch; import com.automq.stream.s3.network.ThrottleStrategy; import com.automq.stream.s3.operator.ObjectStorage; @@ -65,6 +66,7 @@ public class DefaultReader { private final ObjectStorage objectStorage; private final String nodePrefix; private final boolean shardedFirst; + private final S3RequestSource source; private final Time time; private final Queue singleReadTasks = new ConcurrentLinkedQueue<>(); @@ -80,19 +82,20 @@ public class DefaultReader { private final EventLoop eventLoop; public DefaultReader(ObjectStorage objectStorage, String clusterId, int nodeId, String type, Time time) { - this(objectStorage, clusterId, nodeId, type, time, false); + this(objectStorage, clusterId, nodeId, type, time, false, S3RequestSource.UNCLASSIFIED); } /** Creates a dual-layout reader that prefers the configured write layout and falls back only on not-found. */ public DefaultReader(ObjectStorage objectStorage, ObjectWALConfig config, Time time) { - this(objectStorage, config.clusterId(), config.nodeId(), config.type(), time, config.version() == 1); + this(objectStorage, config.clusterId(), config.nodeId(), config.type(), time, config.version() == 1, config.source()); } private DefaultReader(ObjectStorage objectStorage, String clusterId, int nodeId, String type, Time time, - boolean shardedFirst) { + boolean shardedFirst, S3RequestSource source) { this.objectStorage = objectStorage; this.nodePrefix = ObjectUtils.nodePrefix(clusterId, nodeId, type); this.shardedFirst = shardedFirst; + this.source = source; this.time = time; this.eventLoop = EVENT_LOOPS[Math.abs(nodeId % EVENT_LOOPS.length)]; } @@ -280,7 +283,7 @@ private CompletableFuture rangeRead(long epoch, long objectStartOffset, private CompletableFuture rangeRead(String path, long start, long end) { return objectStorage.rangeRead( - new ObjectStorage.ReadOptions().bucket(objectStorage.bucketId()).throttleStrategy(ThrottleStrategy.BYPASS), + new ObjectStorage.ReadOptions().bucket(objectStorage.bucketId()).throttleStrategy(ThrottleStrategy.BYPASS).source(source), path, start, end); } diff --git a/s3stream/src/main/java/com/automq/stream/s3/wal/impl/object/DefaultWriter.java b/s3stream/src/main/java/com/automq/stream/s3/wal/impl/object/DefaultWriter.java index 25091e9ff8..a6a94e9c34 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/wal/impl/object/DefaultWriter.java +++ b/s3stream/src/main/java/com/automq/stream/s3/wal/impl/object/DefaultWriter.java @@ -410,7 +410,7 @@ private void uploadBulk0(Bulk bulk) { int objectLength = objectBuffer.readableBytes(); // Enable fast retry. - ObjectStorage.WriteOptions writeOptions = new ObjectStorage.WriteOptions().enableFastRetry(true); + ObjectStorage.WriteOptions writeOptions = new ObjectStorage.WriteOptions().enableFastRetry(true).source(config.source()); String path = config.version() == 1 ? ObjectUtils.genShardedObjectPathV1(nodePrefix, config.epoch(), firstOffset, endOffset) : ObjectUtils.genObjectPathV1(nodePrefix, config.epoch(), firstOffset, endOffset); diff --git a/s3stream/src/main/java/com/automq/stream/s3/wal/impl/object/ObjectReservationService.java b/s3stream/src/main/java/com/automq/stream/s3/wal/impl/object/ObjectReservationService.java index 94dc0f8979..2b4451f1fd 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/wal/impl/object/ObjectReservationService.java +++ b/s3stream/src/main/java/com/automq/stream/s3/wal/impl/object/ObjectReservationService.java @@ -21,6 +21,7 @@ import com.automq.stream.s3.ByteBufAlloc; import com.automq.stream.s3.Constants; +import com.automq.stream.s3.metrics.operations.S3RequestSource; import com.automq.stream.s3.network.ThrottleStrategy; import com.automq.stream.s3.operator.ObjectStorage; import com.automq.stream.s3.wal.ReservationService; @@ -28,6 +29,7 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import java.util.Objects; import java.util.concurrent.CompletableFuture; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentMap; @@ -47,14 +49,21 @@ public class ObjectReservationService implements ReservationService { private final String clusterId; private final ObjectStorage objectStorage; private final short bucketId; + private final S3RequestSource source; private final ConcurrentMap nodeIdPathMap; public ObjectReservationService(String clusterId, ObjectStorage objectStorage, short bucketId) { + this(clusterId, objectStorage, bucketId, S3RequestSource.UNCLASSIFIED); + } + + /** Attributes ownership-verification GETs to the WAL whose append completion they guard. */ + public ObjectReservationService(String clusterId, ObjectStorage objectStorage, short bucketId, S3RequestSource source) { this.clusterId = clusterId; this.objectStorage = objectStorage; this.nodeIdPathMap = new ConcurrentHashMap<>(); this.bucketId = bucketId; + this.source = Objects.requireNonNull(source, "source"); } private String path(long nodeId) { @@ -63,7 +72,7 @@ private String path(long nodeId) { // Visible for testing protected CompletableFuture verify(long nodeId, ByteBuf target) { - ObjectStorage.ReadOptions options = new ObjectStorage.ReadOptions().throttleStrategy(ThrottleStrategy.BYPASS).bucket(bucketId); + ObjectStorage.ReadOptions options = new ObjectStorage.ReadOptions().throttleStrategy(ThrottleStrategy.BYPASS).bucket(bucketId).source(source); return objectStorage.rangeRead(options, path(nodeId), 0, S3_RESERVATION_OBJECT_LENGTH) .thenApply(bytes -> { try { diff --git a/s3stream/src/main/java/com/automq/stream/s3/wal/impl/object/ObjectWALConfig.java b/s3stream/src/main/java/com/automq/stream/s3/wal/impl/object/ObjectWALConfig.java index 2fb72c0fd9..90a81a12a3 100644 --- a/s3stream/src/main/java/com/automq/stream/s3/wal/impl/object/ObjectWALConfig.java +++ b/s3stream/src/main/java/com/automq/stream/s3/wal/impl/object/ObjectWALConfig.java @@ -19,12 +19,15 @@ package com.automq.stream.s3.wal.impl.object; +import com.automq.stream.s3.metrics.operations.S3RequestSource; import com.automq.stream.s3.wal.OpenMode; import com.automq.stream.s3.wal.ReservationService; import com.automq.stream.utils.IdURI; import org.apache.commons.lang3.StringUtils; +import java.util.Objects; + public class ObjectWALConfig { private final String uri; private final ReservationService reservationService; @@ -39,6 +42,7 @@ public class ObjectWALConfig { private final OpenMode openMode; private final short bucketId; private final String type; + private final S3RequestSource source; private final boolean manualMode; private final int version; @@ -63,6 +67,7 @@ public ObjectWALConfig(String uri, ReservationService reservationService, long b this.openMode = openMode; this.bucketId = bucketId; this.type = type; + this.source = S3RequestSource.UNCLASSIFIED; this.manualMode = false; this.version = 0; } @@ -81,6 +86,7 @@ private ObjectWALConfig(Builder builder) { this.openMode = builder.openMode; this.bucketId = builder.bucketId; this.type = builder.type; + this.source = builder.source; this.manualMode = builder.manualMode; this.version = builder.version; } @@ -143,6 +149,11 @@ public short bucketId() { return bucketId; } + /** Returns the metric ownership of this WAL, independently of its object-key type. */ + public S3RequestSource source() { + return source; + } + public String type() { return type; } @@ -179,6 +190,7 @@ public static final class Builder { private OpenMode openMode = OpenMode.READ_WRITE; private short bucketId; private String type = ""; + private S3RequestSource source = S3RequestSource.UNCLASSIFIED; private boolean manualMode; private int version; @@ -293,6 +305,12 @@ public Builder withBucketId(short bucketId) { return this; } + /** Sets metric ownership for readers and writers; existing callers remain unclassified. */ + public Builder withSource(S3RequestSource source) { + this.source = Objects.requireNonNull(source, "source"); + return this; + } + public Builder withType(String type) { this.type = type; return this; diff --git a/s3stream/src/test/java/com/automq/stream/s3/metrics/MetricsHistogramTest.java b/s3stream/src/test/java/com/automq/stream/s3/metrics/MetricsHistogramTest.java index 39fc2768f7..bebbbe540c 100644 --- a/s3stream/src/test/java/com/automq/stream/s3/metrics/MetricsHistogramTest.java +++ b/s3stream/src/test/java/com/automq/stream/s3/metrics/MetricsHistogramTest.java @@ -9,23 +9,27 @@ import java.util.Queue; import java.util.function.Consumer; +import io.opentelemetry.api.common.AttributeKey; import io.opentelemetry.api.common.Attributes; import io.opentelemetry.api.metrics.DoubleGaugeBuilder; import io.opentelemetry.api.metrics.LongGaugeBuilder; import io.opentelemetry.api.metrics.Meter; import io.opentelemetry.api.metrics.ObservableDoubleGauge; +import io.opentelemetry.api.metrics.ObservableDoubleMeasurement; import io.opentelemetry.api.metrics.ObservableLongGauge; import io.opentelemetry.api.metrics.ObservableLongMeasurement; import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyDouble; import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @Tag("S3Unit") class MetricsHistogramTest { - /** Verifies HistogramBundle count and sum gauges expose cumulative values across interval snapshots. */ + /** Count/sum remain cumulative across snapshots, and all five exports retain the request source. */ @Test void histogramBundleShouldReportCumulativeCountAndSum() throws Exception { Metrics metrics = Metrics.instance(); @@ -35,7 +39,8 @@ void histogramBundleShouldReportCumulativeCountAndSum() throws Exception { setField(metrics, "meter", null); setField(metrics, "globalConfig", null); Metrics.HistogramBundle bundle = metrics.histogram("test_histogram", "test", "nanoseconds"); - DeltaHistogram histogram = bundle.histogram(MetricsLevel.INFO, Attributes.empty()); + Attributes source = Attributes.of(AttributeKey.stringKey("source"), "ROUTER_WAL"); + DeltaHistogram histogram = bundle.histogram(MetricsLevel.INFO, source); Meter meter = mock(Meter.class); DoubleGaugeBuilder doubleGaugeBuilder = mock(DoubleGaugeBuilder.class); @@ -73,10 +78,18 @@ void histogramBundleShouldReportCumulativeCountAndSum() throws Exception { countCallback.accept(countMeasurement); sumCallback.accept(sumMeasurement); - verify(countMeasurement).record(1L, Attributes.empty()); - verify(countMeasurement).record(2L, Attributes.empty()); - verify(sumMeasurement).record(10L, Attributes.empty()); - verify(sumMeasurement).record(30L, Attributes.empty()); + verify(countMeasurement).record(1L, source); + verify(countMeasurement).record(2L, source); + verify(sumMeasurement).record(10L, source); + verify(sumMeasurement).record(30L, source); + for (int i = 0; i < 3; i++) { + @SuppressWarnings("unchecked") + Consumer callback = (Consumer) getInvocationArgument( + doubleGaugeBuilder, i); + ObservableDoubleMeasurement measurement = mock(ObservableDoubleMeasurement.class); + callback.accept(measurement); + verify(measurement).record(anyDouble(), eq(source)); + } } finally { getWaitingSetups(metrics).remove(bundle); setField(metrics, "meter", previousMeter); @@ -84,7 +97,7 @@ void histogramBundleShouldReportCumulativeCountAndSum() throws Exception { } } - private static Object getInvocationArgument(LongGaugeBuilder builder, int invocationIndex) { + private static Object getInvocationArgument(Object builder, int invocationIndex) { return org.mockito.Mockito.mockingDetails(builder) .getInvocations() .stream() diff --git a/s3stream/src/test/java/com/automq/stream/s3/operator/ObjectStorageSourceTest.java b/s3stream/src/test/java/com/automq/stream/s3/operator/ObjectStorageSourceTest.java new file mode 100644 index 0000000000..50efe16c67 --- /dev/null +++ b/s3stream/src/test/java/com/automq/stream/s3/operator/ObjectStorageSourceTest.java @@ -0,0 +1,225 @@ +/* + * Copyright 2026, AutoMQ HK Limited. + * + * 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.operator; + +import com.automq.stream.s3.metrics.operations.S3RequestSource; +import com.automq.stream.s3.metrics.stats.OperationLatencyMetrics; +import com.automq.stream.s3.metrics.wrapper.DeltaHistogram; + +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Tag; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.CsvSource; +import org.junit.jupiter.params.provider.ValueSource; + +import java.lang.reflect.Field; +import java.util.Collection; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.TimeoutException; + +import io.netty.buffer.ByteBuf; +import io.netty.buffer.Unpooled; +import io.opentelemetry.api.common.AttributeKey; +import io.opentelemetry.api.common.Attributes; + +import static org.awaitility.Awaitility.await; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyInt; +import static org.mockito.ArgumentMatchers.anyLong; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.ArgumentMatchers.argThat; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.doReturn; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.spy; +import static org.mockito.Mockito.timeout; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +/** Verifies real GET/PUT completion observations, including merging and speculative retries. */ +@Tag("S3Unit") +@Timeout(10) +class ObjectStorageSourceTest { + private AbstractObjectStorage storage; + + @BeforeEach + void setUp() { + storage = spy(new MemoryObjectStorage(true)); + } + + @AfterEach + void tearDown() { + storage.close(); + } + + /** Normal PUT and GET retries retain ownership and record exactly the existing failure/success deltas. */ + @Test + void normalRetriesPreserveSource() throws Exception { + doReturn(0).when(storage).retryDelay(any(), anyInt()); + doReturn(CompletableFuture.failedFuture(new TimeoutException()), CompletableFuture.completedFuture(null)) + .when(storage).doWrite(any(), anyString(), any()); + long puts = count("put_object", S3RequestSource.DELTA_WAL, "success"); + long failedPuts = count("put_object", S3RequestSource.DELTA_WAL, "failed"); + storage.write(new ObjectStorage.WriteOptions().source(S3RequestSource.DELTA_WAL), "key", data()).get(); + assertEquals(puts + 1, count("put_object", S3RequestSource.DELTA_WAL, "success")); + assertEquals(failedPuts + 1, count("put_object", S3RequestSource.DELTA_WAL, "failed")); + verify(storage, times(2)).doWrite(argThat(o -> o.source() == S3RequestSource.DELTA_WAL), eq("key"), any()); + + doReturn(CompletableFuture.failedFuture(new TimeoutException()), CompletableFuture.completedFuture(data())) + .when(storage).doRangeRead(any(), anyString(), anyLong(), anyLong()); + long gets = count("get_object", S3RequestSource.ROUTER_WAL, "success"); + long failedGets = count("get_object", S3RequestSource.ROUTER_WAL, "failed"); + storage.mergedRangeRead(new ObjectStorage.ReadOptions().source(S3RequestSource.ROUTER_WAL), "key", 0, 64) + .get().release(); + await().untilAsserted(() -> assertEquals(failedGets + 1, count("get_object", S3RequestSource.ROUTER_WAL, "failed"))); + assertEquals(gets + 1, count("get_object", S3RequestSource.ROUTER_WAL, "success")); + verify(storage, times(2)).doRangeRead(argThat(o -> o.source() == S3RequestSource.ROUTER_WAL), eq("key"), eq(0L), eq(64L)); + } + + /** Source differences never split a merge; mutating caller options after submission cannot relabel it. */ + @ParameterizedTest + @CsvSource({"ROUTER_WAL,ROUTER_WAL,ROUTER_WAL", "ROUTER_WAL,DELTA_WAL,MIXED", "ROUTER_WAL,UNCLASSIFIED,MIXED"}) + void mergedReadsHaveOneSource(S3RequestSource first, S3RequestSource second, S3RequestSource expected) throws Exception { + storage.write(ObjectStorage.WriteOptions.DEFAULT, "key", data()).get(); + long before = count("get_object", expected, "success"); + ObjectStorage.ReadOptions options = new ObjectStorage.ReadOptions().source(first); + CompletableFuture left = storage.rangeRead(options, "key", 0, 32); + CompletableFuture right = storage.rangeRead(new ObjectStorage.ReadOptions().source(second), "key", 32, 64); + options.source(S3RequestSource.MIXED); + storage.tryMergeRead(); + left.get().release(); + right.get().release(); + assertEquals(before + 1, count("get_object", expected, "success")); + verify(storage).doRangeRead(argThat(o -> o.source() == expected), eq("key"), eq(0L), eq(64L)); + } + + /** Both successful original and hedge completions count, whichever finishes first, without leaking buffers. */ + @ParameterizedTest + @ValueSource(booleans = {true, false}) + void readHedgePreservesSource(boolean hedgeWins) throws Exception { + setLatency("readLatencyCalculator"); + CompletableFuture original = new CompletableFuture<>(); + CompletableFuture hedge = new CompletableFuture<>(); + doReturn(original, hedge).when(storage).doRangeRead(any(), anyString(), anyLong(), anyLong()); + long before = count("get_object", S3RequestSource.ROUTER_WAL, "success"); + CompletableFuture result = storage.mergedRangeRead( + new ObjectStorage.ReadOptions().source(S3RequestSource.ROUTER_WAL), "key", 0, 64); + verify(storage, timeout(1000).times(2)).doRangeRead(argThat(o -> o.source() == S3RequestSource.ROUTER_WAL), + eq("key"), eq(0L), eq(64L)); + ByteBuf winner = data(); + (hedgeWins ? hedge : original).complete(winner); + result.get().release(); + ByteBuf loser = data(); + (hedgeWins ? original : hedge).complete(loser); + await().untilAsserted(() -> assertEquals(before + 2, count("get_object", S3RequestSource.ROUTER_WAL, "success"))); + assertEquals(0, winner.refCnt()); + assertEquals(0, loser.refCnt()); + } + + /** PUT hedge failures count only while the original is pending; caller mutation cannot change either source. */ + @ParameterizedTest + @ValueSource(booleans = {true, false}) + void writeHedgeFailurePreservesSource(boolean failureBeforeSuccess) throws Exception { + setLatency("s3LatencyCalculator"); + CompletableFuture original = new CompletableFuture<>(); + CompletableFuture hedge = new CompletableFuture<>(); + doReturn(original, hedge).when(storage).doWrite(any(), anyString(), any()); + long succeeded = count("put_object", S3RequestSource.DELTA_WAL, "success"); + long failed = count("put_object", S3RequestSource.DELTA_WAL, "failed"); + ObjectStorage.WriteOptions options = new ObjectStorage.WriteOptions().source(S3RequestSource.DELTA_WAL).enableFastRetry(true); + ByteBuf buffer = data(); + CompletableFuture result = storage.write(options, "key", buffer); + options.source(S3RequestSource.ROUTER_WAL); + verify(storage, timeout(1000).times(2)).doWrite(argThat(o -> o.source() == S3RequestSource.DELTA_WAL), eq("key"), any()); + if (failureBeforeSuccess) { + hedge.completeExceptionally(new TimeoutException()); + assertFalse(result.isDone()); + } + original.complete(null); + result.get(); + if (!failureBeforeSuccess) { + hedge.completeExceptionally(new TimeoutException()); + } + await().untilAsserted(() -> assertEquals(0, buffer.refCnt())); + assertEquals(succeeded + 1, count("put_object", S3RequestSource.DELTA_WAL, "success")); + assertEquals(failed + (failureBeforeSuccess ? 1 : 0), count("put_object", S3RequestSource.DELTA_WAL, "failed")); + } + + /** Unannotated calls are explicit and repeated observations reuse one histogram registration per tuple. */ + @Test + void defaultSourceAndRegistrationAreBounded() throws Exception { + storage.write(ObjectStorage.WriteOptions.DEFAULT, "key", data()).get(); + int registered = histograms().size(); + long before = count("put_object", S3RequestSource.UNCLASSIFIED, "success"); + for (int i = 0; i < 5; i++) { + storage.write(ObjectStorage.WriteOptions.DEFAULT, "key", data()).get(); + } + assertEquals(registered, histograms().size()); + assertEquals(before + 5, count("put_object", S3RequestSource.UNCLASSIFIED, "success")); + assertEquals(S3RequestSource.UNCLASSIFIED, ObjectStorage.WriteOptions.DEFAULT.source()); + } + + private void setLatency(String fieldName) throws Exception { + S3LatencyCalculator calculator = mock(S3LatencyCalculator.class); + when(calculator.valueAtPercentile(anyLong(), anyLong())).thenReturn(10L); + Field field = AbstractObjectStorage.class.getDeclaredField(fieldName); + field.setAccessible(true); + field.set(storage, calculator); + } + + private static ByteBuf data() { + return Unpooled.buffer(64).writeZero(64); + } + + // Inspect the actual registered observations consumed by all five histogram export callbacks. + private static Collection histograms() throws Exception { + Field bundle = OperationLatencyMetrics.class.getDeclaredField("OPERATION_LATENCY"); + bundle.setAccessible(true); + return (Collection) field(bundle.get(null), "histograms"); + } + + private static long count(String operation, S3RequestSource source, String status) throws Exception { + long count = 0; + for (Object registration : histograms()) { + Attributes attributes = (Attributes) field(registration, "histogramAttrs"); + if (operation.equals(attributes.get(AttributeKey.stringKey("operation_name"))) + && source.name().equals(attributes.get(AttributeKey.stringKey("source"))) + && status.equals(attributes.get(AttributeKey.stringKey("status"))) + && "16KB".equals(attributes.get(AttributeKey.stringKey("size")))) { + assertEquals("S3Request", attributes.get(AttributeKey.stringKey("operation_type"))); + assertEquals("16KB", attributes.get(AttributeKey.stringKey("size"))); + count += ((DeltaHistogram) field(registration, "histogram")).cumulativeCount(); + } + } + return count; + } + + private static Object field(Object target, String name) throws Exception { + Field field = target.getClass().getDeclaredField(name); + field.setAccessible(true); + return field.get(target); + } +} diff --git a/s3stream/src/test/java/com/automq/stream/s3/wal/impl/object/ObjectWALSourceTest.java b/s3stream/src/test/java/com/automq/stream/s3/wal/impl/object/ObjectWALSourceTest.java new file mode 100644 index 0000000000..8820c7ccd5 --- /dev/null +++ b/s3stream/src/test/java/com/automq/stream/s3/wal/impl/object/ObjectWALSourceTest.java @@ -0,0 +1,137 @@ +/* + * Copyright 2026, AutoMQ HK Limited. + * + * 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.wal.impl.object; + +import com.automq.stream.s3.DefaultByteBufSupplier; +import com.automq.stream.s3.metrics.operations.S3RequestSource; +import com.automq.stream.s3.model.StreamRecordBatch; +import com.automq.stream.s3.operator.MemoryObjectStorage; +import com.automq.stream.s3.trace.context.TraceContext; +import com.automq.stream.s3.wal.AppendResult; +import com.automq.stream.s3.wal.OpenMode; +import com.automq.stream.utils.Time; + +import org.junit.jupiter.api.Tag; +import org.junit.jupiter.api.Timeout; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.CsvSource; +import org.junit.jupiter.params.provider.ValueSource; + +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.TimeUnit; + +import io.netty.buffer.ByteBuf; +import io.netty.buffer.Unpooled; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyLong; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.ArgumentMatchers.argThat; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.clearInvocations; +import static org.mockito.Mockito.doReturn; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.spy; +import static org.mockito.Mockito.timeout; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; + +/** Covers source ownership at the shared WAL append/read and durability-verification boundaries. */ +@Tag("S3Unit") +@Timeout(10) +class ObjectWALSourceTest { + /** Both WALs use direct PUT and retain ownership on remote reads and both migration fallback layouts. */ + @ParameterizedTest + @CsvSource({"ROUTER_WAL,0", "ROUTER_WAL,1", "DELTA_WAL,0", "DELTA_WAL,1"}) + void appendAndFallbackRead(S3RequestSource source, int version) throws Exception { + MemoryObjectStorage storage = spy(new MemoryObjectStorage()); + ObjectWALConfig config = config(source, version).build(); + ObjectWALService wal = new ObjectWALService(Time.SYSTEM, storage, config); + wal.start(); + try { + CompletableFuture append = wal.append(TraceContext.DEFAULT, record()); + assertFalse(append.isDone()); + wal.flush().get(5, TimeUnit.SECONDS); + AppendResult result = append.get(); + verify(storage).write(argThat(o -> o.source() == source && o.enableFastRetry()), anyString(), any()); + verify(storage, never()).writer(any(), anyString()); + DefaultReader reader = new DefaultReader(storage, + config(source, 1 - version).withOpenMode(OpenMode.READ_ONLY).build(), Time.SYSTEM); + StreamRecordBatch read = reader.get(result.recordOffset()).get(5, TimeUnit.SECONDS); + try { + assertEquals(1, read.getCount()); + } finally { + read.release(); + } + verify(storage, times(2)).rangeRead(argThat(o -> o.source() == source), anyString(), anyLong(), anyLong()); + } finally { + wal.shutdownGracefully(); + storage.close(); + } + } + + /** Delta append completion waits for its labeled ownership GET, including the fenced result. */ + @ParameterizedTest + @ValueSource(booleans = {true, false}) + void appendWaitsForReservation(boolean valid) throws Exception { + MemoryObjectStorage storage = spy(new MemoryObjectStorage()); + ObjectReservationService reservation = new ObjectReservationService("cluster", storage, (short) 0, S3RequestSource.DELTA_WAL); + reservation.acquire(1, 2, false).get(); + ObjectWALService wal = new ObjectWALService(Time.SYSTEM, storage, + config(S3RequestSource.DELTA_WAL, 0).withReservationService(reservation).build()); + wal.start(); + clearInvocations(storage); + CompletableFuture verification = new CompletableFuture<>(); + String path = "reservation/_kafka_cluster/1"; + doReturn(verification).when(storage).rangeRead(any(), eq(path), eq(0L), eq((long) ObjectReservationService.S3_RESERVATION_OBJECT_LENGTH)); + try { + CompletableFuture append = wal.append(TraceContext.DEFAULT, record()); + CompletableFuture flush = wal.flush(); + verify(storage, timeout(1000)).rangeRead(argThat(o -> o.source() == S3RequestSource.DELTA_WAL), + eq(path), eq(0L), eq((long) ObjectReservationService.S3_RESERVATION_OBJECT_LENGTH)); + assertFalse(append.isDone()); + verification.complete(Unpooled.buffer(ObjectReservationService.S3_RESERVATION_OBJECT_LENGTH) + .writeInt(ObjectReservationService.S3_RESERVATION_OBJECT_MAGIC_CODE).writeLong(1) + .writeLong(valid ? 2 : 3).writeBoolean(false)); + if (valid) { + append.get(5, TimeUnit.SECONDS); + flush.get(5, TimeUnit.SECONDS); + } else { + assertThrows(ExecutionException.class, () -> append.get(5, TimeUnit.SECONDS)); + } + } finally { + wal.shutdownGracefully(); + storage.close(); + } + } + + private static ObjectWALConfig.Builder config(S3RequestSource source, int version) { + return ObjectWALConfig.builder().withSource(source).withVersion(version).withClusterId("cluster") + .withNodeId(1).withEpoch(2).withManualMode(true); + } + + private static StreamRecordBatch record() { + return StreamRecordBatch.of(1, 0, 0, 1, Unpooled.buffer(16).writeZero(16), DefaultByteBufSupplier.INSTANCE); + } +}