From f517fdec0a6b118478225a1390b17b816ca47165 Mon Sep 17 00:00:00 2001 From: Hazel Date: Wed, 24 Jun 2026 13:24:47 -0700 Subject: [PATCH 01/65] update pom and config --- pom.xml | 12 +++ .../jsonapi/config/BillingS3ExportConfig.java | 73 +++++++++++++++++++ 2 files changed, 85 insertions(+) create mode 100644 src/main/java/io/stargate/sgv2/jsonapi/config/BillingS3ExportConfig.java diff --git a/pom.xml b/pom.xml index 7239c5a2f6..16e6770ef7 100644 --- a/pom.xml +++ b/pom.xml @@ -203,6 +203,11 @@ software.amazon.awssdk sts + + + software.amazon.awssdk + s3 + com.datastax.oss java-driver-core @@ -300,6 +305,13 @@ junit-jupiter test + + + com.adobe.testing + s3mock-testcontainers + 5.1.0 + test + com.github.docker-java docker-java-api diff --git a/src/main/java/io/stargate/sgv2/jsonapi/config/BillingS3ExportConfig.java b/src/main/java/io/stargate/sgv2/jsonapi/config/BillingS3ExportConfig.java new file mode 100644 index 0000000000..48c98b3a5e --- /dev/null +++ b/src/main/java/io/stargate/sgv2/jsonapi/config/BillingS3ExportConfig.java @@ -0,0 +1,73 @@ +package io.stargate.sgv2.jsonapi.config; + +import io.smallrye.config.ConfigMapping; +import io.smallrye.config.WithDefault; +import java.time.Duration; +import java.util.Optional; + +/** + * Configuration for exporting {@code billing.events} log lines to S3 as NDJSON {@code .jsonl} + * objects. Consumed by {@link io.stargate.sgv2.jsonapi.service.provider.BillingS3HandlerInstaller} + * which, when {@link #enabled()} is {@code true}, attaches a {@link + * io.stargate.sgv2.jsonapi.service.provider.BillingS3LogHandler} to the {@code billing.events} + * logger. + * + *

This is a startup-time switch, not a per-request feature flag — it is independent of + * {@link io.stargate.sgv2.jsonapi.config.feature.ApiFeature#BILLING_EVENTS_LOGGING}. Events only + * reach the handler when the {@code billing.events} logger is also emitting (i.e. the billing + * feature is on); this flag then decides whether those lines are additionally shipped to S3. The + * existing console handler stays attached as a backstop regardless. + * + *

Off by default. When enabled, {@link #bucket()} and {@link #bucketRegion()} are + * required; if either is missing the handler is not installed (logged as an error) and billing + * events continue to flow to the console only. + */ +@ConfigMapping(prefix = "stargate.jsonapi.billing.s3") +public interface BillingS3ExportConfig { + + /** + * Master switch; when {@code false} (default) no handler is installed and no S3 client is built. + */ + @WithDefault("false") + boolean enabled(); + + /** Target bucket, e.g. {@code serverless-usage-dev}. Required when {@link #enabled()}. */ + Optional bucket(); + + /** AWS region of the bucket, e.g. {@code us-east-1}. Required when {@link #enabled()}. */ + Optional bucketRegion(); + + /** + * Endpoint override for the S3 client. Set this to point at a non-AWS S3 (e.g. S3Mock in tests); + * when present, path-style addressing is forced. SDK resolves the regional AWS endpoint when left + * empty. + */ + Optional endpointOverride(); + + /** Seal a batch once it holds this many events. */ + @WithDefault("50") + int maxEvents(); + + /** Seal a batch once its NDJSON body reaches this many bytes (~2 MiB default). */ + @WithDefault("2097152") + long maxBytes(); + + /** Seal an open (under-filled) batch once its oldest event is this old (flush interval). */ + @WithDefault("PT30S") + Duration maxAge(); + + /** + * Capacity of the handler's in-memory hand-off queue. {@link BillingS3LogHandler#publish()} + * offers lines non-blocking; once the queue is full, further lines are dropped and counted. + */ + @WithDefault("10000") + int queueCapacity(); + + /** Maximum number of PUT attempts per sealed batch before it is counted as failed. */ + @WithDefault("3") + int maxUploadAttempts(); + + /** Base delay for exponential backoff between PUT attempts ({@code base * 2^(attempt-1)}). */ + @WithDefault("PT0.2S") + Duration retryBaseBackoff(); +} From 00de470eda6a4bc5a5620f828df42bd7eaf78265 Mon Sep 17 00:00:00 2001 From: Hazel Date: Tue, 30 Jun 2026 15:52:44 -0700 Subject: [PATCH 02/65] add BillingS3LogHandler --- .../service/provider/BillingS3LogHandler.java | 352 ++++++++++++++++++ 1 file changed, 352 insertions(+) create mode 100644 src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandler.java diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandler.java b/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandler.java new file mode 100644 index 0000000000..c732bb4cf2 --- /dev/null +++ b/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandler.java @@ -0,0 +1,352 @@ +package io.stargate.sgv2.jsonapi.service.provider; + +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import io.micrometer.core.instrument.MeterRegistry; +import io.smallrye.mutiny.Multi; +import io.smallrye.mutiny.Uni; +import io.smallrye.mutiny.subscription.BackPressureStrategy; +import io.smallrye.mutiny.subscription.Cancellable; +import io.smallrye.mutiny.subscription.MultiEmitter; +import io.stargate.sgv2.jsonapi.config.BillingS3ExportConfig; +import java.nio.charset.StandardCharsets; +import java.time.Duration; +import java.time.Instant; +import java.time.ZoneOffset; +import java.time.format.DateTimeFormatter; +import java.util.ArrayList; +import java.util.List; +import java.util.UUID; +import java.util.concurrent.CompletionStage; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicLong; +import java.util.logging.Handler; +import java.util.logging.LogRecord; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * A {@link Handler} that ships {@code billing.events} JSON log lines to S3 as NDJSON ({@code + * .jsonl}) objects. Installed on the {@code billing.events} logger by {@link + * BillingS3HandlerInstaller} when {@link BillingS3ExportConfig#enabled()} is {@code true}; the + * existing console handler stays attached as a backstop (dual-write). + * + *

Off the request path. {@link #publish(LogRecord)} only hands the line to an internal + * pipeline — it never blocks and never throws. The pipeline holds a bounded in-memory backlog + * ({@link BillingS3ExportConfig#queueCapacity()}) so transient bursts are absorbed and drained as + * S3 catches up; only when that backlog is full is a line dropped and counted (never silent). The + * pipeline batches lines and seals a batch on {@link BillingS3ExportConfig#maxEvents()} / {@link + * BillingS3ExportConfig#maxBytes()} / {@link BillingS3ExportConfig#maxAge()}, then PUTs it with + * bounded retry/backoff, up to {@link BillingS3ExportConfig#uploadConcurrency()} uploads in flight. + * + *

Verbatim bodies. Each log line is kept byte-for-byte as one NDJSON row — only {@code + * timestamp} is parsed out (for the key's date path); there is no re-serialization. Each sealed + * batch is one object at {@code ///

///.jsonl}; the key is built + * once and reused across retries so a retried PUT overwrites rather than duplicates (downstream + * also dedups on each event id). + * + *

{@link #close()} drains in-flight batches (bounded) and closes the uploader. The handler is + * intentionally not a CDI bean — the installer wires this instance to the {@code + * billing.events} category explicitly. + */ +public final class BillingS3LogHandler extends Handler { + + // ---- Constants ---- + // S3 object-key consistent identifier; TBD + static final String PATH_PREFIX = "billing-events"; + private static final Logger LOG = LoggerFactory.getLogger(BillingS3LogHandler.class); + private static final ObjectMapper MAPPER = new ObjectMapper(); + // UTC, minute-resolution date path for the object key + private static final DateTimeFormatter KEY_TIME_FORMAT = + DateTimeFormatter.ofPattern("yyyy/MM/dd/HH/mm").withZone(ZoneOffset.UTC); + // How long {@link #close()} waits for in-flight batches to drain before cancelling. + private static final long SHUTDOWN_DRAIN_TIMEOUT_MILLIS = 15_000L; + + // ---- Collaborators ---- + private final AsyncBatchUploader uploader; + private final BillingMetrics metrics; + + // ---- Tuning (resolved from BillingS3ExportConfig) ---- + private final int maxEvents; + private final long maxBytes; + private final long maxAgeNanos; + private final int maxUploadAttempts; + private final long retryBaseBackoffMillis; + private final double retryJitter; + private final long queueCapacity; + private final int uploadConcurrency; + + // ---- Reactive pipeline ---- + private volatile MultiEmitter emitter; + private final Cancellable subscription; + private final CountDownLatch terminated = new CountDownLatch(1); + + // ---- Mutable in-flight state ---- + /** + * Events accepted by {@link #publish} but not yet delivered or failed — the in-memory backlog + * depth and the bounded-buffer gate. {@code publish} CAS-checks it against {@link + * #queueCapacity}; {@link BillingMetrics} exposes it as the {@code billing.s3.queue.depth} gauge. + */ + private final AtomicLong backlogEvents = new AtomicLong(0); + + /** + * Current open (unsealed) batch — accumulation state carried across calls. No lock needed: the + * upstream {@code SerializedMultiEmitter} serializes onItem, so {@link #accumulate} and {@link + * #flushOpenBatch} never run concurrently. + */ + private Batch openBatch; + + /** Config-driven constructor used by the installer. */ + public BillingS3LogHandler( + BillingS3ExportConfig config, AsyncBatchUploader uploader, MeterRegistry meterRegistry) { + this( + uploader, + meterRegistry, + config.maxEvents(), + config.maxBytes(), + config.maxAge(), + config.maxUploadAttempts(), + config.retryBaseBackoff(), + config.retryJitter(), + config.queueCapacity(), + config.uploadConcurrency()); + } + + /** Explicit-threshold constructor; convenient for unit tests. */ + BillingS3LogHandler( + AsyncBatchUploader uploader, + MeterRegistry meterRegistry, + int maxEvents, + long maxBytes, + Duration maxAge, + int maxUploadAttempts, + Duration retryBaseBackoff, + double retryJitter, + int queueCapacity, + int uploadConcurrency) { + this.uploader = uploader; + this.maxEvents = Math.max(1, maxEvents); + this.maxBytes = Math.max(1L, maxBytes); + this.maxAgeNanos = Math.max(1L, maxAge.toNanos()); + this.maxUploadAttempts = Math.max(1, maxUploadAttempts); + this.retryBaseBackoffMillis = Math.max(0L, retryBaseBackoff.toMillis()); + this.retryJitter = Math.clamp(retryJitter, 0.0, 1.0); + this.queueCapacity = Math.max(1L, queueCapacity); + this.uploadConcurrency = Math.max(1, uploadConcurrency); + this.metrics = new BillingMetrics(meterRegistry, backlogEvents, this.queueCapacity); + + this.subscription = + Multi.createFrom() + .emitter(em -> this.emitter = em, BackPressureStrategy.BUFFER) + .onItem() + .transform(this::parse) + .onItem() + .transformToMultiAndConcatenate(this::accumulate) + .onCompletion() + .switchTo(this::flushOpenBatch) + .onItem() + .transformToUni(this::uploadWithRetry) + .merge(this.uploadConcurrency) + .onTermination() + .invoke(() -> terminated.countDown()) + .subscribe() + .with( + ignored -> {}, + failure -> LOG.error("Billing S3 export pipeline terminated", failure)); + } + + // ============================================================ + // java.util.logging.Handler + // ============================================================ + + @Override + public void publish(LogRecord record) { + MultiEmitter e = this.emitter; + if (record == null) { + return; + } + String line = record.getMessage(); + if (e == null || line == null || line.isEmpty()) { + return; + } + metrics.recordOffered(); + // Bounded backlog: accept while under capacity (absorbing bursts); once full, shed and count — + // never a silent drop. The CAS keeps the bound exact under concurrent publish(). + long current; + do { + current = backlogEvents.get(); + if (current >= queueCapacity) { + metrics.recordDropped(); + return; + } + } while (!backlogEvents.compareAndSet(current, current + 1)); + e.emit(line); + } + + @Override + public void flush() { + // No-op: the pipeline ships continuously; close() handles the final seal-everything on + // shutdown. + } + + @Override + public void close() { + MultiEmitter e = this.emitter; + if (e != null) { + e.complete(); + } + try { + if (!terminated.await(SHUTDOWN_DRAIN_TIMEOUT_MILLIS, TimeUnit.MILLISECONDS)) { + LOG.warn( + "Billing S3 export did not drain within {} ms on shutdown; cancelling", + SHUTDOWN_DRAIN_TIMEOUT_MILLIS); + } + } catch (InterruptedException ex) { + Thread.currentThread().interrupt(); + } + if (subscription != null) { + subscription.cancel(); + } + try { + uploader.close(); + } catch (Exception ex) { + LOG.warn("Error closing billing S3 uploader", ex); + } + } + + // ============================================================ + // Batching + // ============================================================ + + /** + * Sequential fold of one parsed row into {@link #openBatch}, emitting 0–2 sealed {@link Batch}es + * (sealed on {@code maxEvents}/{@code maxBytes}, or the prior batch on {@code maxAge} at the next + * arrival). Kept ordered and one-at-a-time via {@code …AndConcatenate}, not merge. + */ + private Multi accumulate(Parsed parsed) { + List sealed = new ArrayList<>(2); + Batch batch = openBatch; + if (batch != null && System.nanoTime() - batch.firstNanos >= maxAgeNanos) { + sealed.add(batch); + batch = null; + openBatch = null; + } + if (batch == null) { + batch = new Batch(parsed.timestamp()); + openBatch = batch; + } + batch.add(parsed.line()); + if (batch.events >= maxEvents || batch.bytes >= maxBytes) { + sealed.add(batch); + openBatch = null; + } + return Multi.createFrom().iterable(sealed); + } + + /** Emits the final open batch (if any) when the stream completes. */ + private Multi flushOpenBatch() { + Batch remaining = openBatch; + openBatch = null; + return remaining == null ? Multi.createFrom().empty() : Multi.createFrom().item(remaining); + } + + /** + * Uploads one sealed batch with bounded retry/backoff/jitter. Never propagates failure: a + * giving-up batch is counted and recovered to a no-op so the pipeline stays alive. + */ + private Uni uploadWithRetry(Batch batch) { + String key = objectKey(batch.firstTimestamp, UUID.randomUUID()); + byte[] body = batch.body(); + + Uni put = Uni.createFrom().completionStage(() -> uploader.upload(key, body)); + if (maxUploadAttempts > 1) { + var retry = put.onFailure().retry(); + put = + retryBaseBackoffMillis > 0 + ? retry + .withBackOff(Duration.ofMillis(retryBaseBackoffMillis)) + .withJitter(retryJitter) + .atMost(maxUploadAttempts - 1) + : retry.atMost(maxUploadAttempts - 1); + } + + return put.onItem() + .invoke( + () -> { + metrics.recordBatchDelivered(batch.events); + backlogEvents.addAndGet(-batch.events); + }) + .onFailure() + .invoke( + t -> LOG.error("Giving up on billing S3 batch '{}' ({} events)", key, batch.events, t)) + .onFailure() + .recoverWithItem( + () -> { + metrics.recordBatchFailed(batch.events); + backlogEvents.addAndGet(-batch.events); + return null; + }); + } + + /** Object key: {@code ///

///.jsonl} (UTC). */ + static String objectKey(Instant timestamp, UUID id) { + return PATH_PREFIX + "/" + KEY_TIME_FORMAT.format(timestamp) + "/" + id + ".jsonl"; + } + + /** Parses the {@code timestamp} for the key's date path; keeps the verbatim line. */ + private Parsed parse(String line) { + try { + JsonNode node = MAPPER.readTree(line); + JsonNode tsNode = node.get("timestamp"); + if (tsNode != null && tsNode.isTextual()) { + return new Parsed(line, Instant.parse(tsNode.asText())); + } + } catch (Exception e) { + // fall through to the wall-clock fallback below + } + metrics.recordParseFailure(); + return new Parsed(line, Instant.now()); + } + + private record Parsed(String line, Instant timestamp) {} + + /** A growing set of verbatim NDJSON rows. */ + private static final class Batch { + private final Instant firstTimestamp; + private final long firstNanos = System.nanoTime(); + private final List lines = new ArrayList<>(); + private int events = 0; + private long bytes = 0; + + Batch(Instant firstTimestamp) { + this.firstTimestamp = firstTimestamp; + } + + void add(String line) { + lines.add(line); + events++; + bytes += line.getBytes(StandardCharsets.UTF_8).length + 1L; // +1 for the newline + } + + byte[] body() { + StringBuilder sb = new StringBuilder((int) Math.min(Integer.MAX_VALUE, bytes + events)); + for (String line : lines) { + sb.append(line).append('\n'); + } + return sb.toString().getBytes(StandardCharsets.UTF_8); + } + } + + /** + * Single-attempt async uploader of one sealed batch (test seam); the production implementation is + * {@link S3BatchUploader}. A failed {@link CompletionStage} triggers the handler's retry/backoff. + */ + @FunctionalInterface + public interface AsyncBatchUploader extends AutoCloseable { + CompletionStage upload(String key, byte[] body); + + @Override + default void close() {} + } +} From 4c711c1f17fa7ace60ca16afa0a7ec403d19147e Mon Sep 17 00:00:00 2001 From: Hazel Date: Tue, 30 Jun 2026 15:54:11 -0700 Subject: [PATCH 03/65] update config --- .../sgv2/jsonapi/config/BillingS3ExportConfig.java | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/src/main/java/io/stargate/sgv2/jsonapi/config/BillingS3ExportConfig.java b/src/main/java/io/stargate/sgv2/jsonapi/config/BillingS3ExportConfig.java index 48c98b3a5e..4c242ca9c8 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/config/BillingS3ExportConfig.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/config/BillingS3ExportConfig.java @@ -70,4 +70,12 @@ public interface BillingS3ExportConfig { /** Base delay for exponential backoff between PUT attempts ({@code base * 2^(attempt-1)}). */ @WithDefault("PT0.2S") Duration retryBaseBackoff(); + + /** Jitter factor [0,1] on retry back-off (0 = none). Only applies when retryBaseBackoff > 0. */ + @WithDefault("0.5") + double retryJitter(); + + /** Number of batch uploads (S3 PUTs) allowed in flight concurrently. */ + @WithDefault("4") + int uploadConcurrency(); } From e141bda9b0f432a11afca66c0b435d385fbcfbe7 Mon Sep 17 00:00:00 2001 From: Hazel Date: Wed, 1 Jul 2026 12:42:13 -0700 Subject: [PATCH 04/65] Handler - rename+add comment --- .../service/provider/BillingS3LogHandler.java | 29 +++++++++++++++---- 1 file changed, 24 insertions(+), 5 deletions(-) diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandler.java b/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandler.java index c732bb4cf2..0a10e156ea 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandler.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandler.java @@ -79,7 +79,7 @@ public final class BillingS3LogHandler extends Handler { // ---- Reactive pipeline ---- private volatile MultiEmitter emitter; - private final Cancellable subscription; + private final Cancellable pipeline; private final CountDownLatch terminated = new CountDownLatch(1); // ---- Mutable in-flight state ---- @@ -136,20 +136,38 @@ public BillingS3LogHandler( this.uploadConcurrency = Math.max(1, uploadConcurrency); this.metrics = new BillingMetrics(meterRegistry, backlogEvents, this.queueCapacity); - this.subscription = + // Build + subscribe the export pipeline; runs for the life of the handler. + this.pipeline = Multi.createFrom() + // Source: publish() (from any thread) emits raw JSON lines here — Mutiny's + // SerializedMultiEmitter funnels the concurrent emits into one serial stream, so + // everything downstream runs single-threaded. BUFFER holds the backlog, bounded by the + // publish() capacity gate, so nothing is dropped at this stage. .emitter(em -> this.emitter = em, BackPressureStrategy.BUFFER) .onItem() + // Pull the timestamp for the key's date path; keep the line byte-for-byte, never drop. .transform(this::parse) .onItem() + // Fold each line into the open batch, emitting 0–2 sealed batches (on + // maxEvents/maxBytes, or the prior batch on maxAge). Concatenate -> ordered, + // one-at-a-time .transformToMultiAndConcatenate(this::accumulate) .onCompletion() + // On shutdown (emitter completed), flush the final under-filled batch so nothing is + // stranded. .switchTo(this::flushOpenBatch) .onItem() + // PUT each sealed batch to S3 with bounded retry/backoff… .transformToUni(this::uploadWithRetry) + // …up to uploadConcurrency uploads in flight at once. .merge(this.uploadConcurrency) .onTermination() + // On any terminal (complete/fail/cancel), release the latch close() blocks on for + // graceful drain. .invoke(() -> terminated.countDown()) + // Subscribe -> activates the whole chain. Per-batch result is ignored; only a + // pipeline-fatal failure is logged (uploadWithRetry already recovers per-batch + // failures). .subscribe() .with( ignored -> {}, @@ -167,7 +185,7 @@ public void publish(LogRecord record) { return; } String line = record.getMessage(); - if (e == null || line == null || line.isEmpty()) { + if (e == null || line == null || line.isBlank()) { return; } metrics.recordOffered(); @@ -205,8 +223,8 @@ public void close() { } catch (InterruptedException ex) { Thread.currentThread().interrupt(); } - if (subscription != null) { - subscription.cancel(); + if (pipeline != null) { + pipeline.cancel(); } try { uploader.close(); @@ -260,6 +278,7 @@ private Uni uploadWithRetry(Batch batch) { byte[] body = batch.body(); Uni put = Uni.createFrom().completionStage(() -> uploader.upload(key, body)); + // just retry no if statement if (maxUploadAttempts > 1) { var retry = put.onFailure().retry(); put = From adc2856ce74b406e60d8f1391d86fa89d809b1f5 Mon Sep 17 00:00:00 2001 From: Hazel Date: Mon, 6 Jul 2026 13:00:55 -0700 Subject: [PATCH 05/65] Add S3BatchUploader --- .../service/provider/S3BatchUploader.java | 83 +++++++++++++++++++ 1 file changed, 83 insertions(+) create mode 100644 src/main/java/io/stargate/sgv2/jsonapi/service/provider/S3BatchUploader.java diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/S3BatchUploader.java b/src/main/java/io/stargate/sgv2/jsonapi/service/provider/S3BatchUploader.java new file mode 100644 index 0000000000..a382901b54 --- /dev/null +++ b/src/main/java/io/stargate/sgv2/jsonapi/service/provider/S3BatchUploader.java @@ -0,0 +1,83 @@ +package io.stargate.sgv2.jsonapi.service.provider; + +import java.net.URI; +import java.util.Objects; +import java.util.Optional; +import java.util.concurrent.CompletionStage; +import software.amazon.awssdk.auth.credentials.DefaultCredentialsProvider; +import software.amazon.awssdk.core.async.AsyncRequestBody; +import software.amazon.awssdk.regions.Region; +import software.amazon.awssdk.services.s3.S3AsyncClient; +import software.amazon.awssdk.services.s3.model.PutObjectRequest; + +/** + * Production {@link BillingS3LogHandler.AsyncBatchUploader} backed by an AWS SDK v2 {@link + * S3AsyncClient}: each sealed batch is one async {@code PutObject}, returned as a {@link + * CompletionStage} so the handler's pipeline drives it without blocking. The async client's Netty + * HTTP backend is already on the classpath (pulled by {@code bedrockruntime}), so this needs no new + * sync HTTP-client dependency. + * + *

Credentials come from the {@link DefaultCredentialsProvider} chain (IRSA web-identity token in + * AWS deployments). When an {@code endpointOverride} is configured (e.g. S3Mock in tests), + * path-style addressing is forced so bucket-as-host resolution does not get in the way. + */ +public class S3BatchUploader implements BillingS3LogHandler.AsyncBatchUploader { + + private static final String NDJSON_CONTENT_TYPE = "application/x-ndjson"; + + private final S3AsyncClient client; + private final String bucket; + + S3BatchUploader(S3AsyncClient client, String bucket) { + this.client = client; + this.bucket = bucket; + } + + /** + * Builds an uploader from resolved config. {@code region} and {@code bucket} must be non-null. + */ + public static S3BatchUploader create( + String region, String bucket, Optional endpointOverride) { + Objects.requireNonNull(region, "region must not be null"); + Objects.requireNonNull(bucket, "bucket must not be null"); + + var builder = + S3AsyncClient.builder() + .region(Region.of(region)) + // Credentials resolve from the SDK's default provider chain (env vars, + // web-identity/OIDC token, instance/container roles). This transparently supports + // federated (AssumeRoleWithWebIdentity) and cross-account access — the bucket may live + // in a different account (per IAM + bucket policy); its region is set via .region(). + .credentialsProvider(DefaultCredentialsProvider.create()); + + // Real AWS S3 needs no endpoint: the SDK endpoint rules (s3 SDK's DefaultS3EndpointProvider) + // derive https://.s3..amazonaws.com from region + partition dnsSuffix. + // An override is only for a non-AWS S3 (S3Mock in tests): it bypasses those rules and forces + // path-style, since a localhost host can't virtual-host the bucket as a subdomain. + endpointOverride + .filter(s -> !s.isBlank()) + .ifPresent(uri -> builder.endpointOverride(URI.create(uri)).forcePathStyle(true)); + + return new S3BatchUploader(builder.build(), bucket); + } + + @Override + public CompletionStage upload(String key, byte[] body) { + // Returns the async PUT future (a failed future drives the handler's retry/backoff); no + // blocking. + return client + .putObject( + PutObjectRequest.builder() + .bucket(bucket) + .key(key) + .contentType(NDJSON_CONTENT_TYPE) + .build(), + AsyncRequestBody.fromBytes(body)) + .thenAccept(resp -> {}); + } + + @Override + public void close() { + client.close(); + } +} From a8405fb4d60f98cc9ed901687b47a981d03a807a Mon Sep 17 00:00:00 2001 From: Hazel Date: Mon, 6 Jul 2026 13:02:16 -0700 Subject: [PATCH 06/65] Update BillingS3ExportConfig --- .../io/stargate/sgv2/jsonapi/config/BillingS3ExportConfig.java | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/main/java/io/stargate/sgv2/jsonapi/config/BillingS3ExportConfig.java b/src/main/java/io/stargate/sgv2/jsonapi/config/BillingS3ExportConfig.java index 4c242ca9c8..5f535fc486 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/config/BillingS3ExportConfig.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/config/BillingS3ExportConfig.java @@ -39,8 +39,7 @@ public interface BillingS3ExportConfig { /** * Endpoint override for the S3 client. Set this to point at a non-AWS S3 (e.g. S3Mock in tests); - * when present, path-style addressing is forced. SDK resolves the regional AWS endpoint when left - * empty. + * SDK resolves the regional AWS endpoint when left empty. */ Optional endpointOverride(); From 91720163404ebc32e200e798944c0e03fa7de4fa Mon Sep 17 00:00:00 2001 From: Hazel Date: Mon, 6 Jul 2026 14:23:17 -0700 Subject: [PATCH 07/65] Update retry --- .../service/provider/BillingS3LogHandler.java | 54 +++++++++---------- 1 file changed, 27 insertions(+), 27 deletions(-) diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandler.java b/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandler.java index 0a10e156ea..5c2cca7c98 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandler.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandler.java @@ -72,7 +72,8 @@ public final class BillingS3LogHandler extends Handler { private final long maxBytes; private final long maxAgeNanos; private final int maxUploadAttempts; - private final long retryBaseBackoffMillis; + private final Duration initialBackOffMillis; + private final Duration maxBackOffMillis; private final double retryJitter; private final long queueCapacity; private final int uploadConcurrency; @@ -106,8 +107,9 @@ public BillingS3LogHandler( config.maxEvents(), config.maxBytes(), config.maxAge(), - config.maxUploadAttempts(), - config.retryBaseBackoff(), + config.atMostRetries(), + Duration.ofMillis(config.initialBackOffMillis()), + Duration.ofMillis(config.maxBackOffMillis()), config.retryJitter(), config.queueCapacity(), config.uploadConcurrency()); @@ -121,19 +123,21 @@ public BillingS3LogHandler( long maxBytes, Duration maxAge, int maxUploadAttempts, - Duration retryBaseBackoff, + Duration initialBackOffMillis, + Duration maxBackOffMillis, double retryJitter, int queueCapacity, int uploadConcurrency) { this.uploader = uploader; - this.maxEvents = Math.max(1, maxEvents); - this.maxBytes = Math.max(1L, maxBytes); - this.maxAgeNanos = Math.max(1L, maxAge.toNanos()); - this.maxUploadAttempts = Math.max(1, maxUploadAttempts); - this.retryBaseBackoffMillis = Math.max(0L, retryBaseBackoff.toMillis()); - this.retryJitter = Math.clamp(retryJitter, 0.0, 1.0); - this.queueCapacity = Math.max(1L, queueCapacity); - this.uploadConcurrency = Math.max(1, uploadConcurrency); + this.maxEvents = maxEvents; + this.maxBytes = maxBytes; + this.maxAgeNanos = maxAge.toNanos(); + this.maxUploadAttempts = maxUploadAttempts; + this.initialBackOffMillis = initialBackOffMillis; + this.maxBackOffMillis = maxBackOffMillis; + this.retryJitter = retryJitter; + this.queueCapacity = queueCapacity; + this.uploadConcurrency = uploadConcurrency; this.metrics = new BillingMetrics(meterRegistry, backlogEvents, this.queueCapacity); // Build + subscribe the export pipeline; runs for the life of the handler. @@ -277,20 +281,14 @@ private Uni uploadWithRetry(Batch batch) { String key = objectKey(batch.firstTimestamp, UUID.randomUUID()); byte[] body = batch.body(); - Uni put = Uni.createFrom().completionStage(() -> uploader.upload(key, body)); - // just retry no if statement - if (maxUploadAttempts > 1) { - var retry = put.onFailure().retry(); - put = - retryBaseBackoffMillis > 0 - ? retry - .withBackOff(Duration.ofMillis(retryBaseBackoffMillis)) - .withJitter(retryJitter) - .atMost(maxUploadAttempts - 1) - : retry.atMost(maxUploadAttempts - 1); - } - - return put.onItem() + return Uni.createFrom() + .completionStage(() -> uploader.upload(key, body)) + .onFailure() + .retry() + .withBackOff(initialBackOffMillis, maxBackOffMillis) + .withJitter(retryJitter) + .atMost(maxUploadAttempts - 1) + .onItem() .invoke( () -> { metrics.recordBatchDelivered(batch.events); @@ -298,7 +296,9 @@ private Uni uploadWithRetry(Batch batch) { }) .onFailure() .invoke( - t -> LOG.error("Giving up on billing S3 batch '{}' ({} events)", key, batch.events, t)) + t -> + LOG.error( + "Failed to upload billing S3 batch '{}' ({} events)", key, batch.events, t)) .onFailure() .recoverWithItem( () -> { From 1a605142d986b649b5b5cb5365f3f58ba8245574 Mon Sep 17 00:00:00 2001 From: Hazel Date: Mon, 6 Jul 2026 14:23:29 -0700 Subject: [PATCH 08/65] Update config --- .../jsonapi/config/BillingS3ExportConfig.java | 24 +++++++++++++++---- 1 file changed, 19 insertions(+), 5 deletions(-) diff --git a/src/main/java/io/stargate/sgv2/jsonapi/config/BillingS3ExportConfig.java b/src/main/java/io/stargate/sgv2/jsonapi/config/BillingS3ExportConfig.java index 5f535fc486..6cd3f38c94 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/config/BillingS3ExportConfig.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/config/BillingS3ExportConfig.java @@ -62,13 +62,27 @@ public interface BillingS3ExportConfig { @WithDefault("10000") int queueCapacity(); - /** Maximum number of PUT attempts per sealed batch before it is counted as failed. */ + /** + * Maximum number of PUT attempts per sealed batch before it is counted as failed. Default is 3 (1 + * request + 2 retries). + */ @WithDefault("3") - int maxUploadAttempts(); + int atMostRetries(); + + // /** Base delay for exponential backoff between PUT attempts ({@code base * 2^(attempt-1)}). */ + // @WithDefault("PT0.2S") + // Duration retryBaseBackoff(); + + /** + * The initial delay between retries in milliseconds. The first retry occurs after the specified + * delay (default 100 ms), doubling each time until reaching maxBackOffMillis. + */ + @WithDefault("100") + int initialBackOffMillis(); - /** Base delay for exponential backoff between PUT attempts ({@code base * 2^(attempt-1)}). */ - @WithDefault("PT0.2S") - Duration retryBaseBackoff(); + /** The maximum delay between retries in milliseconds. */ + @WithDefault("500") + int maxBackOffMillis(); /** Jitter factor [0,1] on retry back-off (0 = none). Only applies when retryBaseBackoff > 0. */ @WithDefault("0.5") From fb8bcf928b8c5ad2087d399432e06be61bbd58f7 Mon Sep 17 00:00:00 2001 From: Hazel Date: Mon, 6 Jul 2026 17:13:09 -0700 Subject: [PATCH 09/65] Update config --- .../jsonapi/config/BillingS3ExportConfig.java | 12 ++++------ .../service/provider/BillingS3LogHandler.java | 22 +++++++++---------- 2 files changed, 15 insertions(+), 19 deletions(-) diff --git a/src/main/java/io/stargate/sgv2/jsonapi/config/BillingS3ExportConfig.java b/src/main/java/io/stargate/sgv2/jsonapi/config/BillingS3ExportConfig.java index 6cd3f38c94..f1f910dd66 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/config/BillingS3ExportConfig.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/config/BillingS3ExportConfig.java @@ -63,16 +63,12 @@ public interface BillingS3ExportConfig { int queueCapacity(); /** - * Maximum number of PUT attempts per sealed batch before it is counted as failed. Default is 3 (1 - * request + 2 retries). + * Maximum number of retries after a failed PUT, per sealed batch, before the batch is counted as + * failed. Default is 2 retries (up to 3 attempts including the initial PUT). */ - @WithDefault("3") + @WithDefault("2") int atMostRetries(); - // /** Base delay for exponential backoff between PUT attempts ({@code base * 2^(attempt-1)}). */ - // @WithDefault("PT0.2S") - // Duration retryBaseBackoff(); - /** * The initial delay between retries in milliseconds. The first retry occurs after the specified * delay (default 100 ms), doubling each time until reaching maxBackOffMillis. @@ -84,7 +80,7 @@ public interface BillingS3ExportConfig { @WithDefault("500") int maxBackOffMillis(); - /** Jitter factor [0,1] on retry back-off (0 = none). Only applies when retryBaseBackoff > 0. */ + /** A random variation added to the delay between retries in an exponential backoff strategy. */ @WithDefault("0.5") double retryJitter(); diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandler.java b/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandler.java index 5c2cca7c98..bfbcdda497 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandler.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandler.java @@ -71,9 +71,9 @@ public final class BillingS3LogHandler extends Handler { private final int maxEvents; private final long maxBytes; private final long maxAgeNanos; - private final int maxUploadAttempts; - private final Duration initialBackOffMillis; - private final Duration maxBackOffMillis; + private final int atMostRetries; + private final Duration initialBackOffDuration; + private final Duration maxBackOffDuration; private final double retryJitter; private final long queueCapacity; private final int uploadConcurrency; @@ -122,9 +122,9 @@ public BillingS3LogHandler( int maxEvents, long maxBytes, Duration maxAge, - int maxUploadAttempts, - Duration initialBackOffMillis, - Duration maxBackOffMillis, + int atMostRetries, + Duration initialBackOffDuration, + Duration maxBackOffDuration, double retryJitter, int queueCapacity, int uploadConcurrency) { @@ -132,9 +132,9 @@ public BillingS3LogHandler( this.maxEvents = maxEvents; this.maxBytes = maxBytes; this.maxAgeNanos = maxAge.toNanos(); - this.maxUploadAttempts = maxUploadAttempts; - this.initialBackOffMillis = initialBackOffMillis; - this.maxBackOffMillis = maxBackOffMillis; + this.atMostRetries = atMostRetries; + this.initialBackOffDuration = initialBackOffDuration; + this.maxBackOffDuration = maxBackOffDuration; this.retryJitter = retryJitter; this.queueCapacity = queueCapacity; this.uploadConcurrency = uploadConcurrency; @@ -285,9 +285,9 @@ private Uni uploadWithRetry(Batch batch) { .completionStage(() -> uploader.upload(key, body)) .onFailure() .retry() - .withBackOff(initialBackOffMillis, maxBackOffMillis) + .withBackOff(initialBackOffDuration, maxBackOffDuration) .withJitter(retryJitter) - .atMost(maxUploadAttempts - 1) + .atMost(atMostRetries) .onItem() .invoke( () -> { From 8da34f43f3ca047213f0ed405e0befa839933804 Mon Sep 17 00:00:00 2001 From: Hazel Date: Mon, 6 Jul 2026 18:06:11 -0700 Subject: [PATCH 10/65] Add BillingS3HandlerInstaller --- .../provider/BillingS3HandlerInstaller.java | 105 ++++++++++++++++++ 1 file changed, 105 insertions(+) create mode 100644 src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3HandlerInstaller.java diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3HandlerInstaller.java b/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3HandlerInstaller.java new file mode 100644 index 0000000000..69c4e0c577 --- /dev/null +++ b/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3HandlerInstaller.java @@ -0,0 +1,105 @@ +package io.stargate.sgv2.jsonapi.service.provider; + +import io.micrometer.core.instrument.MeterRegistry; +import io.quarkus.runtime.ShutdownEvent; +import io.quarkus.runtime.StartupEvent; +import io.stargate.sgv2.jsonapi.config.BillingS3ExportConfig; +import jakarta.enterprise.context.ApplicationScoped; +import jakarta.enterprise.event.Observes; +import jakarta.inject.Inject; +import java.util.logging.Logger; +import org.slf4j.LoggerFactory; + +/** + * Wires a {@link BillingS3LogHandler} onto the {@code billing.events} logger at startup when {@link + * BillingS3ExportConfig#enabled()} is {@code true}, and removes/closes it on shutdown for a + * graceful drain. + * + *

We attach the handler directly to the {@code billing.events} JUL logger rather than relying on + * Quarkus's discovered-{@code Handler}-bean mechanism: discovered handler beans are attached to the + * root logger, but {@code billing.events} is configured {@code use-parent-handlers: false} + * (so it would never feed a root handler) and we want this handler scoped to exactly that category. + * Adding the handler here in a {@link StartupEvent} observer runs after Quarkus has applied its + * logging configuration, so the registration sticks; we keep a strong reference to the logger so it + * (and our handler) cannot be GC'd. + * + *

Injecting {@link BillingS3ExportConfig} also pins it as a SmallRye {@code @ConfigMapping} bean + * so Quarkus ARC does not drop it at build time (see {@link + * io.stargate.sgv2.jsonapi.JsonApiStartUp} for the same pattern with {@code BillingConfig}). + */ +@ApplicationScoped +public class BillingS3HandlerInstaller { + + private static final org.slf4j.Logger LOG = + LoggerFactory.getLogger(BillingS3HandlerInstaller.class); + + static final String BILLING_LOGGER_NAME = "billing.events"; + + private final BillingS3ExportConfig config; + private final MeterRegistry meterRegistry; + + // Strong references so the configured logger (and the handler we add to it) are not collected, + // and so we can detach cleanly on shutdown. + private volatile Logger billingLogger; + private volatile BillingS3LogHandler handler; + + @Inject + public BillingS3HandlerInstaller(BillingS3ExportConfig config, MeterRegistry meterRegistry) { + this.config = config; + this.meterRegistry = meterRegistry; + } + + void onStart(@Observes StartupEvent event) { + if (!config.enabled()) { + LOG.debug("Billing S3 export disabled (stargate.jsonapi.billing.s3.enabled=false)"); + return; + } + + String bucket = config.bucket().filter(s -> !s.isBlank()).orElse(null); + String region = config.bucketRegion().filter(s -> !s.isBlank()).orElse(null); + if (bucket == null || region == null) { + LOG.error( + "Billing S3 export is enabled but bucket/region are not fully configured (bucket={}," + + " region={}); handler NOT installed. Billing events continue to the console only.", + config.bucket().orElse(""), + config.bucketRegion().orElse("")); + return; + } + + try { + S3BatchUploader uploader = S3BatchUploader.create(region, bucket, config.endpointOverride()); + BillingS3LogHandler newHandler = new BillingS3LogHandler(config, uploader, meterRegistry); + + Logger logger = Logger.getLogger(BILLING_LOGGER_NAME); + logger.addHandler(newHandler); + + this.handler = newHandler; + this.billingLogger = logger; + LOG.info( + "Installed billing S3 export handler on '{}' → bucket '{}' (region '{}', endpointOverride={})", + BILLING_LOGGER_NAME, + bucket, + region, + config.endpointOverride().orElse("")); + } catch (Exception e) { + LOG.error( + "Failed to install billing S3 export handler; billing events continue to the console only", + e); + } + } + + void onStop(@Observes ShutdownEvent event) { + BillingS3LogHandler current = this.handler; + if (current == null) { + return; + } + if (billingLogger != null) { + billingLogger.removeHandler(current); + } + try { + current.close(); // drains remaining batches + } catch (Exception e) { + LOG.warn("Error during billing S3 export handler shutdown", e); + } + } +} From 948c4e1d923db3bfb31f13b912290cc76ac737a6 Mon Sep 17 00:00:00 2001 From: Hazel Date: Mon, 6 Jul 2026 18:24:48 -0700 Subject: [PATCH 11/65] update comments --- .../provider/BillingS3HandlerInstaller.java | 27 ++++++++----------- 1 file changed, 11 insertions(+), 16 deletions(-) diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3HandlerInstaller.java b/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3HandlerInstaller.java index 69c4e0c577..9699e0450f 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3HandlerInstaller.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3HandlerInstaller.java @@ -11,21 +11,17 @@ import org.slf4j.LoggerFactory; /** - * Wires a {@link BillingS3LogHandler} onto the {@code billing.events} logger at startup when {@link - * BillingS3ExportConfig#enabled()} is {@code true}, and removes/closes it on shutdown for a - * graceful drain. + * Attaches a {@link BillingS3LogHandler} to the {@code billing.events} JUL logger at startup (when + * {@link BillingS3ExportConfig#enabled()} is {@code true}) and removes + closes it on shutdown for + * a graceful drain. * - *

We attach the handler directly to the {@code billing.events} JUL logger rather than relying on - * Quarkus's discovered-{@code Handler}-bean mechanism: discovered handler beans are attached to the - * root logger, but {@code billing.events} is configured {@code use-parent-handlers: false} - * (so it would never feed a root handler) and we want this handler scoped to exactly that category. - * Adding the handler here in a {@link StartupEvent} observer runs after Quarkus has applied its - * logging configuration, so the registration sticks; we keep a strong reference to the logger so it - * (and our handler) cannot be GC'd. - * - *

Injecting {@link BillingS3ExportConfig} also pins it as a SmallRye {@code @ConfigMapping} bean - * so Quarkus ARC does not drop it at build time (see {@link - * io.stargate.sgv2.jsonapi.JsonApiStartUp} for the same pattern with {@code BillingConfig}). + *

Done programmatically because Quarkus config can't express it: a category's {@code handlers} + * list can only reference Quarkus's built-in handler types (console/file/syslog/socket), not a + * custom {@link java.util.logging.Handler} class. The one config-driven alternative — a discovered + * {@code @Produces Handler} bean — attaches to the root logger, but {@code billing.events} + * is {@code use-parent-handlers: false} and we want delivery scoped to exactly that category. The + * {@link StartupEvent} observer runs after Quarkus has applied its logging config, so the + * registration sticks. */ @ApplicationScoped public class BillingS3HandlerInstaller { @@ -38,8 +34,7 @@ public class BillingS3HandlerInstaller { private final BillingS3ExportConfig config; private final MeterRegistry meterRegistry; - // Strong references so the configured logger (and the handler we add to it) are not collected, - // and so we can detach cleanly on shutdown. + // Held only to detach + close the handler on shutdown private volatile Logger billingLogger; private volatile BillingS3LogHandler handler; From 888bd0227ddcda88035086d390eaa6d8bb7c24aa Mon Sep 17 00:00:00 2001 From: Hazel Date: Tue, 7 Jul 2026 10:18:32 -0700 Subject: [PATCH 12/65] update installer --- .../provider/BillingS3HandlerInstaller.java | 31 +++++++------------ 1 file changed, 11 insertions(+), 20 deletions(-) diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3HandlerInstaller.java b/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3HandlerInstaller.java index 9699e0450f..a09b85b9ab 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3HandlerInstaller.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3HandlerInstaller.java @@ -34,8 +34,6 @@ public class BillingS3HandlerInstaller { private final BillingS3ExportConfig config; private final MeterRegistry meterRegistry; - // Held only to detach + close the handler on shutdown - private volatile Logger billingLogger; private volatile BillingS3LogHandler handler; @Inject @@ -50,32 +48,28 @@ void onStart(@Observes StartupEvent event) { return; } - String bucket = config.bucket().filter(s -> !s.isBlank()).orElse(null); - String region = config.bucketRegion().filter(s -> !s.isBlank()).orElse(null); + var bucket = config.bucket().filter(s -> !s.isBlank()).orElse(null); + var region = config.bucketRegion().filter(s -> !s.isBlank()).orElse(null); if (bucket == null || region == null) { LOG.error( "Billing S3 export is enabled but bucket/region are not fully configured (bucket={}," + " region={}); handler NOT installed. Billing events continue to the console only.", - config.bucket().orElse(""), - config.bucketRegion().orElse("")); + bucket, + region); return; } try { - S3BatchUploader uploader = S3BatchUploader.create(region, bucket, config.endpointOverride()); - BillingS3LogHandler newHandler = new BillingS3LogHandler(config, uploader, meterRegistry); - - Logger logger = Logger.getLogger(BILLING_LOGGER_NAME); - logger.addHandler(newHandler); + var uploader = S3BatchUploader.create(region, bucket, config.endpointOverride()); + this.handler = new BillingS3LogHandler(config, uploader, meterRegistry); + Logger.getLogger(BILLING_LOGGER_NAME).addHandler(this.handler); - this.handler = newHandler; - this.billingLogger = logger; LOG.info( "Installed billing S3 export handler on '{}' → bucket '{}' (region '{}', endpointOverride={})", BILLING_LOGGER_NAME, bucket, region, - config.endpointOverride().orElse("")); + config.endpointOverride().orElse(null)); } catch (Exception e) { LOG.error( "Failed to install billing S3 export handler; billing events continue to the console only", @@ -84,15 +78,12 @@ void onStart(@Observes StartupEvent event) { } void onStop(@Observes ShutdownEvent event) { - BillingS3LogHandler current = this.handler; - if (current == null) { + if (this.handler == null){ return; } - if (billingLogger != null) { - billingLogger.removeHandler(current); - } + Logger.getLogger(BILLING_LOGGER_NAME).removeHandler(this.handler); try { - current.close(); // drains remaining batches + this.handler.close(); } catch (Exception e) { LOG.warn("Error during billing S3 export handler shutdown", e); } From 758cd0c907ff5d1f0f3b68b1edee1755ffadc680 Mon Sep 17 00:00:00 2001 From: Hazel Date: Tue, 7 Jul 2026 13:22:36 -0700 Subject: [PATCH 13/65] update S3BatchUploader and BillingS3HandlerInstaller --- .../provider/BillingS3HandlerInstaller.java | 46 +++++------ .../service/provider/S3BatchUploader.java | 82 +++++++++++++------ 2 files changed, 78 insertions(+), 50 deletions(-) diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3HandlerInstaller.java b/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3HandlerInstaller.java index a09b85b9ab..66865738e3 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3HandlerInstaller.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3HandlerInstaller.java @@ -7,6 +7,7 @@ import jakarta.enterprise.context.ApplicationScoped; import jakarta.enterprise.event.Observes; import jakarta.inject.Inject; +import java.time.Duration; import java.util.logging.Logger; import org.slf4j.LoggerFactory; @@ -48,37 +49,30 @@ void onStart(@Observes StartupEvent event) { return; } - var bucket = config.bucket().filter(s -> !s.isBlank()).orElse(null); - var region = config.bucketRegion().filter(s -> !s.isBlank()).orElse(null); - if (bucket == null || region == null) { - LOG.error( - "Billing S3 export is enabled but bucket/region are not fully configured (bucket={}," - + " region={}); handler NOT installed. Billing events continue to the console only.", - bucket, - region); - return; - } + var uploader = + S3BatchUploader.create( + config.bucketRegion().orElse(null), + config.bucket().orElse(null), + config.endpointOverride(), + new S3BatchUploader.RetryPolicy( + config.atMostRetries(), + Duration.ofMillis(config.initialBackOffMillis()), + Duration.ofMillis(config.maxBackOffMillis()), + config.retryJitter())); - try { - var uploader = S3BatchUploader.create(region, bucket, config.endpointOverride()); - this.handler = new BillingS3LogHandler(config, uploader, meterRegistry); - Logger.getLogger(BILLING_LOGGER_NAME).addHandler(this.handler); + this.handler = new BillingS3LogHandler(config, uploader, meterRegistry); + Logger.getLogger(BILLING_LOGGER_NAME).addHandler(this.handler); - LOG.info( - "Installed billing S3 export handler on '{}' → bucket '{}' (region '{}', endpointOverride={})", - BILLING_LOGGER_NAME, - bucket, - region, - config.endpointOverride().orElse(null)); - } catch (Exception e) { - LOG.error( - "Failed to install billing S3 export handler; billing events continue to the console only", - e); - } + LOG.info( + "Installed billing S3 export handler on '{}' → bucket '{}' (region '{}', endpointOverride={})", + BILLING_LOGGER_NAME, + config.bucket(), + config.bucketRegion(), + config.endpointOverride().orElse(null)); } void onStop(@Observes ShutdownEvent event) { - if (this.handler == null){ + if (this.handler == null) { return; } Logger.getLogger(BILLING_LOGGER_NAME).removeHandler(this.handler); diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/S3BatchUploader.java b/src/main/java/io/stargate/sgv2/jsonapi/service/provider/S3BatchUploader.java index a382901b54..95f39971d5 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/S3BatchUploader.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/service/provider/S3BatchUploader.java @@ -1,9 +1,10 @@ package io.stargate.sgv2.jsonapi.service.provider; +import io.smallrye.mutiny.Uni; import java.net.URI; +import java.time.Duration; import java.util.Objects; import java.util.Optional; -import java.util.concurrent.CompletionStage; import software.amazon.awssdk.auth.credentials.DefaultCredentialsProvider; import software.amazon.awssdk.core.async.AsyncRequestBody; import software.amazon.awssdk.regions.Region; @@ -11,11 +12,9 @@ import software.amazon.awssdk.services.s3.model.PutObjectRequest; /** - * Production {@link BillingS3LogHandler.AsyncBatchUploader} backed by an AWS SDK v2 {@link - * S3AsyncClient}: each sealed batch is one async {@code PutObject}, returned as a {@link - * CompletionStage} so the handler's pipeline drives it without blocking. The async client's Netty - * HTTP backend is already on the classpath (pulled by {@code bedrockruntime}), so this needs no new - * sync HTTP-client dependency. + * {@link BillingS3LogHandler.AsyncBatchUploader} backed by an AWS SDK v2 {@link S3AsyncClient}: + * each sealed batch is one async {@code PutObject} with bounded retry/backoff, returned as a {@link + * Uni} so the handler's pipeline drives it without blocking. * *

Credentials come from the {@link DefaultCredentialsProvider} chain (IRSA web-identity token in * AWS deployments). When an {@code endpointOverride} is configured (e.g. S3Mock in tests), @@ -27,19 +26,26 @@ public class S3BatchUploader implements BillingS3LogHandler.AsyncBatchUploader { private final S3AsyncClient client; private final String bucket; + private final RetryPolicy retry; - S3BatchUploader(S3AsyncClient client, String bucket) { + S3BatchUploader(S3AsyncClient client, String bucket, RetryPolicy retry) { this.client = client; this.bucket = bucket; + this.retry = retry; } /** - * Builds an uploader from resolved config. {@code region} and {@code bucket} must be non-null. + * Builds an uploader from resolved inputs. {@code endpointOverride} is present only for a non-AWS + * S3 (e.g. S3Mock in tests). */ public static S3BatchUploader create( - String region, String bucket, Optional endpointOverride) { - Objects.requireNonNull(region, "region must not be null"); - Objects.requireNonNull(bucket, "bucket must not be null"); + String region, String bucket, Optional endpointOverride, RetryPolicy retry) { + if (region == null || region.isBlank()) + throw new IllegalArgumentException("stargate.jsonapi.billing.s3.bucket-region must be set"); + if (bucket == null || bucket.isBlank()) + throw new IllegalArgumentException("stargate.jsonapi.billing.s3.bucket must be set"); + Objects.requireNonNull(endpointOverride, "endpointOverride must not be null"); + Objects.requireNonNull(retry, "retry must not be null"); var builder = S3AsyncClient.builder() @@ -58,26 +64,54 @@ public static S3BatchUploader create( .filter(s -> !s.isBlank()) .ifPresent(uri -> builder.endpointOverride(URI.create(uri)).forcePathStyle(true)); - return new S3BatchUploader(builder.build(), bucket); + return new S3BatchUploader(builder.build(), bucket, retry); } @Override - public CompletionStage upload(String key, byte[] body) { - // Returns the async PUT future (a failed future drives the handler's retry/backoff); no - // blocking. - return client - .putObject( - PutObjectRequest.builder() - .bucket(bucket) - .key(key) - .contentType(NDJSON_CONTENT_TYPE) - .build(), - AsyncRequestBody.fromBytes(body)) - .thenAccept(resp -> {}); + public Uni upload(String key, byte[] body) { + return Uni.createFrom() + .completionStage( + () -> + client + .putObject( + PutObjectRequest.builder() + .bucket(bucket) + .key(key) + .contentType(NDJSON_CONTENT_TYPE) + .build(), + AsyncRequestBody.fromBytes(body)) + .thenAccept(resp -> {})) + .onFailure() + .retry() + .withBackOff(retry.initialBackOff(), retry.maxBackOff()) + .withJitter(retry.jitter()) + .atMost(retry.atMostRetries()); } @Override public void close() { client.close(); } + + /** + * Bounded exponential-backoff-with-jitter tuning for one PUT's retries. Validated on construction + * so bad tuning fails at wiring time with a clear message. + */ + public record RetryPolicy( + int atMostRetries, Duration initialBackOff, Duration maxBackOff, double jitter) { + public RetryPolicy { + if (atMostRetries < 0) { + throw new IllegalArgumentException("atMostRetries must be >= 0"); + } + if (initialBackOff.isNegative() || initialBackOff.isZero()) { + throw new IllegalArgumentException("initialBackOff must be > 0"); + } + if (maxBackOff.compareTo(initialBackOff) < 0) { + throw new IllegalArgumentException("maxBackOff must be >= initialBackOff"); + } + if (jitter < 0 || jitter > 1) { + throw new IllegalArgumentException("jitter must be in [0, 1]"); + } + } + } } From 51e95494bc59033b22338823187669c114da76fb Mon Sep 17 00:00:00 2001 From: Hazel Date: Tue, 7 Jul 2026 13:44:48 -0700 Subject: [PATCH 14/65] update BillingS3HandlerInstaller --- .../service/provider/BillingS3HandlerInstaller.java | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3HandlerInstaller.java b/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3HandlerInstaller.java index 66865738e3..b912a45809 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3HandlerInstaller.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3HandlerInstaller.java @@ -49,10 +49,14 @@ void onStart(@Observes StartupEvent event) { return; } + var region = config.bucketRegion().orElse(null); + var bucket = config.bucket().orElse(null); + + // Fail-loud: invalid billing S3 config throws here, aborting application startup. var uploader = S3BatchUploader.create( - config.bucketRegion().orElse(null), - config.bucket().orElse(null), + region, + bucket, config.endpointOverride(), new S3BatchUploader.RetryPolicy( config.atMostRetries(), @@ -66,8 +70,8 @@ void onStart(@Observes StartupEvent event) { LOG.info( "Installed billing S3 export handler on '{}' → bucket '{}' (region '{}', endpointOverride={})", BILLING_LOGGER_NAME, - config.bucket(), - config.bucketRegion(), + bucket, + region, config.endpointOverride().orElse(null)); } From caef7ee21d951abf6e01e769573323a656f765fe Mon Sep 17 00:00:00 2001 From: Hazel Date: Tue, 7 Jul 2026 15:48:20 -0700 Subject: [PATCH 15/65] remove retry --- .../service/provider/BillingS3LogHandler.java | 46 +++++-------------- 1 file changed, 12 insertions(+), 34 deletions(-) diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandler.java b/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandler.java index bfbcdda497..da74702126 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandler.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandler.java @@ -17,7 +17,6 @@ import java.util.ArrayList; import java.util.List; import java.util.UUID; -import java.util.concurrent.CompletionStage; import java.util.concurrent.CountDownLatch; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicLong; @@ -71,10 +70,6 @@ public final class BillingS3LogHandler extends Handler { private final int maxEvents; private final long maxBytes; private final long maxAgeNanos; - private final int atMostRetries; - private final Duration initialBackOffDuration; - private final Duration maxBackOffDuration; - private final double retryJitter; private final long queueCapacity; private final int uploadConcurrency; @@ -107,10 +102,6 @@ public BillingS3LogHandler( config.maxEvents(), config.maxBytes(), config.maxAge(), - config.atMostRetries(), - Duration.ofMillis(config.initialBackOffMillis()), - Duration.ofMillis(config.maxBackOffMillis()), - config.retryJitter(), config.queueCapacity(), config.uploadConcurrency()); } @@ -122,20 +113,12 @@ public BillingS3LogHandler( int maxEvents, long maxBytes, Duration maxAge, - int atMostRetries, - Duration initialBackOffDuration, - Duration maxBackOffDuration, - double retryJitter, int queueCapacity, int uploadConcurrency) { this.uploader = uploader; this.maxEvents = maxEvents; this.maxBytes = maxBytes; this.maxAgeNanos = maxAge.toNanos(); - this.atMostRetries = atMostRetries; - this.initialBackOffDuration = initialBackOffDuration; - this.maxBackOffDuration = maxBackOffDuration; - this.retryJitter = retryJitter; this.queueCapacity = queueCapacity; this.uploadConcurrency = uploadConcurrency; this.metrics = new BillingMetrics(meterRegistry, backlogEvents, this.queueCapacity); @@ -161,8 +144,8 @@ public BillingS3LogHandler( // stranded. .switchTo(this::flushOpenBatch) .onItem() - // PUT each sealed batch to S3 with bounded retry/backoff… - .transformToUni(this::uploadWithRetry) + // PUT each sealed batch to S3 (the uploader applies bounded retry/backoff)… + .transformToUni(this::uploadBatch) // …up to uploadConcurrency uploads in flight at once. .merge(this.uploadConcurrency) .onTermination() @@ -170,8 +153,7 @@ public BillingS3LogHandler( // graceful drain. .invoke(() -> terminated.countDown()) // Subscribe -> activates the whole chain. Per-batch result is ignored; only a - // pipeline-fatal failure is logged (uploadWithRetry already recovers per-batch - // failures). + // pipeline-fatal failure is logged (uploadBatch already recovers per-batch failures). .subscribe() .with( ignored -> {}, @@ -274,20 +256,16 @@ private Multi flushOpenBatch() { } /** - * Uploads one sealed batch with bounded retry/backoff/jitter. Never propagates failure: a - * giving-up batch is counted and recovered to a no-op so the pipeline stays alive. + * Ships one sealed batch: the uploader PUTs it, applying its own bounded retry/backoff. Never + * propagates failure — a giving-up batch is counted and recovered to a no-op so the pipeline + * stays alive. */ - private Uni uploadWithRetry(Batch batch) { + private Uni uploadBatch(Batch batch) { String key = objectKey(batch.firstTimestamp, UUID.randomUUID()); byte[] body = batch.body(); - return Uni.createFrom() - .completionStage(() -> uploader.upload(key, body)) - .onFailure() - .retry() - .withBackOff(initialBackOffDuration, maxBackOffDuration) - .withJitter(retryJitter) - .atMost(atMostRetries) + return uploader + .upload(key, body) .onItem() .invoke( () -> { @@ -358,12 +336,12 @@ byte[] body() { } /** - * Single-attempt async uploader of one sealed batch (test seam); the production implementation is - * {@link S3BatchUploader}. A failed {@link CompletionStage} triggers the handler's retry/backoff. + * Uploads one sealed batch to S3 with its own bounded retry/backoff, failing only once retries + * are exhausted (see implementation {@link S3BatchUploader}). */ @FunctionalInterface public interface AsyncBatchUploader extends AutoCloseable { - CompletionStage upload(String key, byte[] body); + Uni upload(String key, byte[] body); @Override default void close() {} From a7a0f700c37577bffb62d2f967f0b19ec873071f Mon Sep 17 00:00:00 2001 From: Hazel Date: Tue, 7 Jul 2026 16:47:28 -0700 Subject: [PATCH 16/65] Update accumulate method --- .../service/provider/BillingS3LogHandler.java | 24 ++++++++++--------- 1 file changed, 13 insertions(+), 11 deletions(-) diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandler.java b/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandler.java index da74702126..87565209bd 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandler.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandler.java @@ -226,23 +226,25 @@ public void close() { /** * Sequential fold of one parsed row into {@link #openBatch}, emitting 0–2 sealed {@link Batch}es * (sealed on {@code maxEvents}/{@code maxBytes}, or the prior batch on {@code maxAge} at the next - * arrival). Kept ordered and one-at-a-time via {@code …AndConcatenate}, not merge. + * arrival). */ private Multi accumulate(Parsed parsed) { List sealed = new ArrayList<>(2); - Batch batch = openBatch; - if (batch != null && System.nanoTime() - batch.firstNanos >= maxAgeNanos) { - sealed.add(batch); - batch = null; + + // Seal the open batch first if it has aged out, so this line starts a fresh one. + if (openBatch != null && System.nanoTime() - openBatch.firstNanos >= maxAgeNanos) { + sealed.add(openBatch); openBatch = null; } - if (batch == null) { - batch = new Batch(parsed.timestamp()); - openBatch = batch; + if (openBatch == null) { + openBatch = new Batch(parsed.timestamp()); } - batch.add(parsed.line()); - if (batch.events >= maxEvents || batch.bytes >= maxBytes) { - sealed.add(batch); + + openBatch.add(parsed.line()); + + // Seal once the batch is full by count or size. + if (openBatch.events >= maxEvents || openBatch.bytes >= maxBytes) { + sealed.add(openBatch); openBatch = null; } return Multi.createFrom().iterable(sealed); From 109b5d9a99770c90b44ddc207cdb3686d36ec476 Mon Sep 17 00:00:00 2001 From: Hazel Date: Mon, 13 Jul 2026 14:59:58 -0700 Subject: [PATCH 17/65] Update handler --- .../service/provider/BillingS3LogHandler.java | 319 +++++++----------- 1 file changed, 121 insertions(+), 198 deletions(-) diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandler.java b/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandler.java index 87565209bd..9818d614bf 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandler.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandler.java @@ -2,12 +2,10 @@ import com.fasterxml.jackson.databind.JsonNode; import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.common.annotations.VisibleForTesting; import io.micrometer.core.instrument.MeterRegistry; -import io.smallrye.mutiny.Multi; import io.smallrye.mutiny.Uni; -import io.smallrye.mutiny.subscription.BackPressureStrategy; -import io.smallrye.mutiny.subscription.Cancellable; -import io.smallrye.mutiny.subscription.MultiEmitter; +import io.smallrye.mutiny.infrastructure.Infrastructure; import io.stargate.sgv2.jsonapi.config.BillingS3ExportConfig; import java.nio.charset.StandardCharsets; import java.time.Duration; @@ -17,9 +15,11 @@ import java.util.ArrayList; import java.util.List; import java.util.UUID; -import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ArrayBlockingQueue; +import java.util.concurrent.BlockingQueue; import java.util.concurrent.TimeUnit; -import java.util.concurrent.atomic.AtomicLong; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.locks.LockSupport; import java.util.logging.Handler; import java.util.logging.LogRecord; import org.slf4j.Logger; @@ -31,23 +31,22 @@ * BillingS3HandlerInstaller} when {@link BillingS3ExportConfig#enabled()} is {@code true}; the * existing console handler stays attached as a backstop (dual-write). * - *

Off the request path. {@link #publish(LogRecord)} only hands the line to an internal - * pipeline — it never blocks and never throws. The pipeline holds a bounded in-memory backlog - * ({@link BillingS3ExportConfig#queueCapacity()}) so transient bursts are absorbed and drained as - * S3 catches up; only when that backlog is full is a line dropped and counted (never silent). The - * pipeline batches lines and seals a batch on {@link BillingS3ExportConfig#maxEvents()} / {@link - * BillingS3ExportConfig#maxBytes()} / {@link BillingS3ExportConfig#maxAge()}, then PUTs it with - * bounded retry/backoff, up to {@link BillingS3ExportConfig#uploadConcurrency()} uploads in flight. + *

Off the request path. {@link #publish(LogRecord)} only offers the line to a bounded + * in-memory queue — it never blocks and never throws. When the queue holds a full batch ({@link + * BillingS3ExportConfig#maxEvents()} lines) a flush is dispatched to a worker-pool thread: it + * drains up to that many lines and PUTs them as one object. Up to {@link + * BillingS3ExportConfig#uploadConcurrency()} flushes run at once — a non-blocking in-flight counter + * is the gate, and the S3 PUT is async so worker threads are not held during upload. When the queue + * is full, further lines are dropped and counted (never silent). * - *

Verbatim bodies. Each log line is kept byte-for-byte as one NDJSON row — only {@code - * timestamp} is parsed out (for the key's date path); there is no re-serialization. Each sealed - * batch is one object at {@code ///

///.jsonl}; the key is built - * once and reused across retries so a retried PUT overwrites rather than duplicates (downstream - * also dedups on each event id). + *

Verbatim bodies. Each log line is kept byte-for-byte as one NDJSON row; only the + * batch's first line is parsed (for the key's date path). Each flushed batch is one object at + * {@code ///

///.jsonl}; the key is built once and reused across + * the uploader's retries so a retried PUT overwrites rather than duplicates. * - *

{@link #close()} drains in-flight batches (bounded) and closes the uploader. The handler is - * intentionally not a CDI bean — the installer wires this instance to the {@code - * billing.events} category explicitly. + *

{@link #close()} drains whatever remains (including a final under-batch tail), waits for + * in-flight flushes to settle, then closes the uploader. The handler is intentionally not a + * CDI bean — the installer wires this instance to the {@code billing.events} category explicitly. */ public final class BillingS3LogHandler extends Handler { @@ -59,7 +58,7 @@ public final class BillingS3LogHandler extends Handler { // UTC, minute-resolution date path for the object key private static final DateTimeFormatter KEY_TIME_FORMAT = DateTimeFormatter.ofPattern("yyyy/MM/dd/HH/mm").withZone(ZoneOffset.UTC); - // How long {@link #close()} waits for in-flight batches to drain before cancelling. + // How long {@link #close()} waits for the remaining lines and in-flight flushes to drain. private static final long SHUTDOWN_DRAIN_TIMEOUT_MILLIS = 15_000L; // ---- Collaborators ---- @@ -67,31 +66,17 @@ public final class BillingS3LogHandler extends Handler { private final BillingMetrics metrics; // ---- Tuning (resolved from BillingS3ExportConfig) ---- - private final int maxEvents; - private final long maxBytes; - private final long maxAgeNanos; - private final long queueCapacity; - private final int uploadConcurrency; + private final int batchSize; // flush trigger + max lines per object (config.maxEvents) + private final int uploadConcurrency; // max flushes (S3 PUTs) in flight at once - // ---- Reactive pipeline ---- - private volatile MultiEmitter emitter; - private final Cancellable pipeline; - private final CountDownLatch terminated = new CountDownLatch(1); + // ---- Queue + concurrency gate ---- + private final BlockingQueue queue; - // ---- Mutable in-flight state ---- /** - * Events accepted by {@link #publish} but not yet delivered or failed — the in-memory backlog - * depth and the bounded-buffer gate. {@code publish} CAS-checks it against {@link - * #queueCapacity}; {@link BillingMetrics} exposes it as the {@code billing.s3.queue.depth} gauge. + * Flushes (S3 PUTs) currently in flight — the non-blocking concurrency gate, capped at {@link + * #uploadConcurrency}. A slot is held from {@link #drain} through the PUT's completion. */ - private final AtomicLong backlogEvents = new AtomicLong(0); - - /** - * Current open (unsealed) batch — accumulation state carried across calls. No lock needed: the - * upstream {@code SerializedMultiEmitter} serializes onItem, so {@link #accumulate} and {@link - * #flushOpenBatch} never run concurrently. - */ - private Batch openBatch; + private final AtomicInteger inFlight = new AtomicInteger(0); /** Config-driven constructor used by the installer. */ public BillingS3LogHandler( @@ -106,7 +91,12 @@ public BillingS3LogHandler( config.uploadConcurrency()); } - /** Explicit-threshold constructor; convenient for unit tests. */ + /** + * Explicit-threshold constructor; convenient for unit tests. {@code maxBytes} and {@code maxAge} + * are accepted for config compatibility but not yet wired (object byte-cap and time-based flush + * are a follow-up); flushing is currently count-based on {@code maxEvents}. + */ + @VisibleForTesting BillingS3LogHandler( AsyncBatchUploader uploader, MeterRegistry meterRegistry, @@ -116,48 +106,10 @@ public BillingS3LogHandler( int queueCapacity, int uploadConcurrency) { this.uploader = uploader; - this.maxEvents = maxEvents; - this.maxBytes = maxBytes; - this.maxAgeNanos = maxAge.toNanos(); - this.queueCapacity = queueCapacity; + this.batchSize = maxEvents; this.uploadConcurrency = uploadConcurrency; - this.metrics = new BillingMetrics(meterRegistry, backlogEvents, this.queueCapacity); - - // Build + subscribe the export pipeline; runs for the life of the handler. - this.pipeline = - Multi.createFrom() - // Source: publish() (from any thread) emits raw JSON lines here — Mutiny's - // SerializedMultiEmitter funnels the concurrent emits into one serial stream, so - // everything downstream runs single-threaded. BUFFER holds the backlog, bounded by the - // publish() capacity gate, so nothing is dropped at this stage. - .emitter(em -> this.emitter = em, BackPressureStrategy.BUFFER) - .onItem() - // Pull the timestamp for the key's date path; keep the line byte-for-byte, never drop. - .transform(this::parse) - .onItem() - // Fold each line into the open batch, emitting 0–2 sealed batches (on - // maxEvents/maxBytes, or the prior batch on maxAge). Concatenate -> ordered, - // one-at-a-time - .transformToMultiAndConcatenate(this::accumulate) - .onCompletion() - // On shutdown (emitter completed), flush the final under-filled batch so nothing is - // stranded. - .switchTo(this::flushOpenBatch) - .onItem() - // PUT each sealed batch to S3 (the uploader applies bounded retry/backoff)… - .transformToUni(this::uploadBatch) - // …up to uploadConcurrency uploads in flight at once. - .merge(this.uploadConcurrency) - .onTermination() - // On any terminal (complete/fail/cancel), release the latch close() blocks on for - // graceful drain. - .invoke(() -> terminated.countDown()) - // Subscribe -> activates the whole chain. Per-batch result is ignored; only a - // pipeline-fatal failure is logged (uploadBatch already recovers per-batch failures). - .subscribe() - .with( - ignored -> {}, - failure -> LOG.error("Billing S3 export pipeline terminated", failure)); + this.queue = new ArrayBlockingQueue<>(queueCapacity); + this.metrics = new BillingMetrics(meterRegistry, queue::size, queueCapacity); } // ============================================================ @@ -166,124 +118,111 @@ public BillingS3LogHandler( @Override public void publish(LogRecord record) { - MultiEmitter e = this.emitter; if (record == null) { return; } String line = record.getMessage(); - if (e == null || line == null || line.isBlank()) { + if (line == null || line.isBlank()) { return; } metrics.recordOffered(); - // Bounded backlog: accept while under capacity (absorbing bursts); once full, shed and count — - // never a silent drop. The CAS keeps the bound exact under concurrent publish(). - long current; - do { - current = backlogEvents.get(); - if (current >= queueCapacity) { - metrics.recordDropped(); - return; - } - } while (!backlogEvents.compareAndSet(current, current + 1)); - e.emit(line); + if (!queue.offer(line)) { + // Bounded queue full: shed and count — never block, never throw, never silent. + metrics.recordDropped(); + return; + } + maybeFlush(); } @Override public void flush() { - // No-op: the pipeline ships continuously; close() handles the final seal-everything on - // shutdown. + // No-op: flushes are size-triggered and continuous; close() seals whatever remains on shutdown. } + /** + * Shutdown drain: serially flushes the buffered backlog (incl. the sub-batch tail the size-gate + * skips), then waits for in-flight PUTs to finish before closing the client. No internal timeout + * — bounded by the platform's shutdown grace (SIGKILL). + */ @Override public void close() { - MultiEmitter e = this.emitter; - if (e != null) { - e.complete(); + List batch; + while (!(batch = drain()).isEmpty()) { + uploadBatch(batch).await().indefinitely(); } - try { - if (!terminated.await(SHUTDOWN_DRAIN_TIMEOUT_MILLIS, TimeUnit.MILLISECONDS)) { - LOG.warn( - "Billing S3 export did not drain within {} ms on shutdown; cancelling", - SHUTDOWN_DRAIN_TIMEOUT_MILLIS); - } - } catch (InterruptedException ex) { - Thread.currentThread().interrupt(); - } - if (pipeline != null) { - pipeline.cancel(); - } - try { - uploader.close(); - } catch (Exception ex) { - LOG.warn("Error closing billing S3 uploader", ex); + while (inFlight.get() > 0) { + LockSupport.parkNanos(TimeUnit.MILLISECONDS.toNanos(10)); } + uploader.close(); } // ============================================================ - // Batching + // Flush pipeline // ============================================================ /** - * Sequential fold of one parsed row into {@link #openBatch}, emitting 0–2 sealed {@link Batch}es - * (sealed on {@code maxEvents}/{@code maxBytes}, or the prior batch on {@code maxAge} at the next - * arrival). + * Dispatches one flush if the queue holds a full batch and a concurrency slot is free. Called + * after every {@link #publish} and again when a flush completes, so the pipeline self-clocks up + * to {@link #uploadConcurrency} concurrent uploads with no standing reader thread. */ - private Multi accumulate(Parsed parsed) { - List sealed = new ArrayList<>(2); - - // Seal the open batch first if it has aged out, so this line starts a fresh one. - if (openBatch != null && System.nanoTime() - openBatch.firstNanos >= maxAgeNanos) { - sealed.add(openBatch); - openBatch = null; + private void maybeFlush() { + if (queue.size() >= batchSize && tryFlush()) { + Uni.createFrom() + .item(this::drain) + .runSubscriptionOn(Infrastructure.getDefaultWorkerPool()) + .flatMap(this::uploadBatch) + .eventually( + () -> { + inFlight.getAndDecrement(); + maybeFlush(); // a full batch may have accumulated while this one uploaded + }) + .subscribe() + .with(ignored -> {}, failure -> LOG.error("Billing S3 export flush failed", failure)); } - if (openBatch == null) { - openBatch = new Batch(parsed.timestamp()); - } - - openBatch.add(parsed.line()); - - // Seal once the batch is full by count or size. - if (openBatch.events >= maxEvents || openBatch.bytes >= maxBytes) { - sealed.add(openBatch); - openBatch = null; - } - return Multi.createFrom().iterable(sealed); } - /** Emits the final open batch (if any) when the stream completes. */ - private Multi flushOpenBatch() { - Batch remaining = openBatch; - openBatch = null; - return remaining == null ? Multi.createFrom().empty() : Multi.createFrom().item(remaining); + /** + * Non-blocking CAS gate: claims an in-flight slot iff fewer than {@link #uploadConcurrency} are + * held. + */ + private boolean tryFlush() { + int prev = inFlight.getAndUpdate(n -> n < uploadConcurrency ? n + 1 : n); + return prev < uploadConcurrency; } /** - * Ships one sealed batch: the uploader PUTs it, applying its own bounded retry/backoff. Never - * propagates failure — a giving-up batch is counted and recovered to a no-op so the pipeline - * stays alive. + * Removes up to {@link #batchSize} lines from the queue (may come back empty if another flush + * raced ahead and took them first). */ - private Uni uploadBatch(Batch batch) { - String key = objectKey(batch.firstTimestamp, UUID.randomUUID()); - byte[] body = batch.body(); + private List drain() { + List batch = new ArrayList<>(batchSize); + queue.drainTo(batch, batchSize); + return batch; + } + /** + * Ships one drained batch: the uploader PUTs it as a single NDJSON object, applying its own + * bounded retry/backoff. Never propagates failure — a giving-up batch is counted and recovered to + * a no-op so the flush loop stays alive. + */ + private Uni uploadBatch(List batch) { + if (batch.isEmpty()) { + return Uni.createFrom().voidItem(); + } + int events = batch.size(); + String key = objectKey(firstTimestamp(batch), UUID.randomUUID()); + byte[] body = toNdjson(batch); return uploader .upload(key, body) .onItem() - .invoke( - () -> { - metrics.recordBatchDelivered(batch.events); - backlogEvents.addAndGet(-batch.events); - }) + .invoke(() -> metrics.recordBatchDelivered(events)) .onFailure() .invoke( - t -> - LOG.error( - "Failed to upload billing S3 batch '{}' ({} events)", key, batch.events, t)) + t -> LOG.error("Failed to upload billing S3 batch '{}' ({} events)", key, events, t)) .onFailure() .recoverWithItem( () -> { - metrics.recordBatchFailed(batch.events); - backlogEvents.addAndGet(-batch.events); + metrics.recordBatchFailed(events); return null; }); } @@ -293,52 +232,36 @@ static String objectKey(Instant timestamp, UUID id) { return PATH_PREFIX + "/" + KEY_TIME_FORMAT.format(timestamp) + "/" + id + ".jsonl"; } - /** Parses the {@code timestamp} for the key's date path; keeps the verbatim line. */ - private Parsed parse(String line) { + /** NDJSON body: each line verbatim, newline-terminated. */ + private static byte[] toNdjson(List batch) { + StringBuilder sb = new StringBuilder(); + for (String line : batch) { + sb.append(line).append('\n'); + } + return sb.toString().getBytes(StandardCharsets.UTF_8); + } + + /** + * Timestamp for the object key's date path — parsed from the batch's first line's {@code + * timestamp}, falling back to wall clock (and counting the parse failure) when it is + * absent/unparseable. + */ + private Instant firstTimestamp(List batch) { try { - JsonNode node = MAPPER.readTree(line); + JsonNode node = MAPPER.readTree(batch.get(0)); JsonNode tsNode = node.get("timestamp"); if (tsNode != null && tsNode.isTextual()) { - return new Parsed(line, Instant.parse(tsNode.asText())); + return Instant.parse(tsNode.asText()); } } catch (Exception e) { // fall through to the wall-clock fallback below } metrics.recordParseFailure(); - return new Parsed(line, Instant.now()); - } - - private record Parsed(String line, Instant timestamp) {} - - /** A growing set of verbatim NDJSON rows. */ - private static final class Batch { - private final Instant firstTimestamp; - private final long firstNanos = System.nanoTime(); - private final List lines = new ArrayList<>(); - private int events = 0; - private long bytes = 0; - - Batch(Instant firstTimestamp) { - this.firstTimestamp = firstTimestamp; - } - - void add(String line) { - lines.add(line); - events++; - bytes += line.getBytes(StandardCharsets.UTF_8).length + 1L; // +1 for the newline - } - - byte[] body() { - StringBuilder sb = new StringBuilder((int) Math.min(Integer.MAX_VALUE, bytes + events)); - for (String line : lines) { - sb.append(line).append('\n'); - } - return sb.toString().getBytes(StandardCharsets.UTF_8); - } + return Instant.now(); } /** - * Uploads one sealed batch to S3 with its own bounded retry/backoff, failing only once retries + * Uploads one drained batch to S3 with its own bounded retry/backoff, failing only once retries * are exhausted (see implementation {@link S3BatchUploader}). */ @FunctionalInterface From a4e0d9fa3c1f1fc13d734bd6995633b990f04078 Mon Sep 17 00:00:00 2001 From: Hazel Date: Tue, 14 Jul 2026 12:06:56 -0700 Subject: [PATCH 18/65] Add bytes gate --- .../service/provider/BillingS3LogHandler.java | 55 ++++++++++++++----- 1 file changed, 42 insertions(+), 13 deletions(-) diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandler.java b/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandler.java index 9818d614bf..2678633d5d 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandler.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandler.java @@ -19,6 +19,7 @@ import java.util.concurrent.BlockingQueue; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicLong; import java.util.concurrent.locks.LockSupport; import java.util.logging.Handler; import java.util.logging.LogRecord; @@ -58,8 +59,6 @@ public final class BillingS3LogHandler extends Handler { // UTC, minute-resolution date path for the object key private static final DateTimeFormatter KEY_TIME_FORMAT = DateTimeFormatter.ofPattern("yyyy/MM/dd/HH/mm").withZone(ZoneOffset.UTC); - // How long {@link #close()} waits for the remaining lines and in-flight flushes to drain. - private static final long SHUTDOWN_DRAIN_TIMEOUT_MILLIS = 15_000L; // ---- Collaborators ---- private final AsyncBatchUploader uploader; @@ -67,11 +66,18 @@ public final class BillingS3LogHandler extends Handler { // ---- Tuning (resolved from BillingS3ExportConfig) ---- private final int batchSize; // flush trigger + max lines per object (config.maxEvents) + private final long maxBytes; // max NDJSON body bytes per object (config.maxBytes) private final int uploadConcurrency; // max flushes (S3 PUTs) in flight at once - // ---- Queue + concurrency gate ---- + // ---- Queue + byte/concurrency gates ---- private final BlockingQueue queue; + /** + * Running total of NDJSON body bytes buffered in {@link #queue} (each line's length plus its + * newline) — drives the byte-based flush trigger. + */ + private final AtomicLong queuedBytes = new AtomicLong(0); + /** * Flushes (S3 PUTs) currently in flight — the non-blocking concurrency gate, capped at {@link * #uploadConcurrency}. A slot is held from {@link #drain} through the PUT's completion. @@ -92,9 +98,10 @@ public BillingS3LogHandler( } /** - * Explicit-threshold constructor; convenient for unit tests. {@code maxBytes} and {@code maxAge} - * are accepted for config compatibility but not yet wired (object byte-cap and time-based flush - * are a follow-up); flushing is currently count-based on {@code maxEvents}. + * Explicit-threshold constructor; convenient for unit tests. {@code maxAge} is accepted for + * config compatibility but not yet wired (time-based flush is a follow-up); a batch is sealed on + * whichever of {@code maxEvents} (line count) or {@code maxBytes} (NDJSON body size) it hits + * first. */ @VisibleForTesting BillingS3LogHandler( @@ -107,6 +114,7 @@ public BillingS3LogHandler( int uploadConcurrency) { this.uploader = uploader; this.batchSize = maxEvents; + this.maxBytes = maxBytes; this.uploadConcurrency = uploadConcurrency; this.queue = new ArrayBlockingQueue<>(queueCapacity); this.metrics = new BillingMetrics(meterRegistry, queue::size, queueCapacity); @@ -131,6 +139,7 @@ public void publish(LogRecord record) { metrics.recordDropped(); return; } + queuedBytes.addAndGet(lineBytes(line)); maybeFlush(); } @@ -161,12 +170,13 @@ public void close() { // ============================================================ /** - * Dispatches one flush if the queue holds a full batch and a concurrency slot is free. Called - * after every {@link #publish} and again when a flush completes, so the pipeline self-clocks up - * to {@link #uploadConcurrency} concurrent uploads with no standing reader thread. + * Dispatches one flush if a full object's worth is buffered ({@link #shouldFlush}) and a + * concurrency slot is free. Called after every {@link #publish} and again when a flush completes, + * so the pipeline self-clocks up to {@link #uploadConcurrency} concurrent uploads with no + * standing reader thread. */ private void maybeFlush() { - if (queue.size() >= batchSize && tryFlush()) { + if (shouldFlush() && tryFlush()) { Uni.createFrom() .item(this::drain) .runSubscriptionOn(Infrastructure.getDefaultWorkerPool()) @@ -181,6 +191,11 @@ private void maybeFlush() { } } + /** True when the queue holds a full object's worth — by line count or by buffered bytes. */ + private boolean shouldFlush() { + return queue.size() >= batchSize || queuedBytes.get() >= maxBytes; + } + /** * Non-blocking CAS gate: claims an in-flight slot iff fewer than {@link #uploadConcurrency} are * held. @@ -191,15 +206,29 @@ private boolean tryFlush() { } /** - * Removes up to {@link #batchSize} lines from the queue (may come back empty if another flush - * raced ahead and took them first). + * Removes lines from the queue for one object — up to {@link #batchSize} lines or {@link + * #maxBytes} of NDJSON body, whichever comes first (a single over-cap line still goes out alone). + * May come back empty if another flush raced ahead and took them first. */ private List drain() { List batch = new ArrayList<>(batchSize); - queue.drainTo(batch, batchSize); + long bytes = 0; + String line; + while (batch.size() < batchSize && bytes < maxBytes && (line = queue.poll()) != null) { + batch.add(line); + bytes += lineBytes(line); + } + queuedBytes.addAndGet(-bytes); return batch; } + /** + * NDJSON body bytes one line contributes: its length (== UTF-8 bytes for ASCII) plus a newline. + */ + private static int lineBytes(String line) { + return line.length() + 1; + } + /** * Ships one drained batch: the uploader PUTs it as a single NDJSON object, applying its own * bounded retry/backoff. Never propagates failure — a giving-up batch is counted and recovered to From 4f2abee764780e1a71ff4b2af0ca2d495d30f6a2 Mon Sep 17 00:00:00 2001 From: Hazel Date: Tue, 14 Jul 2026 19:21:01 -0700 Subject: [PATCH 19/65] Update S3BatchUploader --- .../service/provider/S3BatchUploader.java | 100 ++++++++---------- 1 file changed, 46 insertions(+), 54 deletions(-) diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/S3BatchUploader.java b/src/main/java/io/stargate/sgv2/jsonapi/service/provider/S3BatchUploader.java index 95f39971d5..31532bb7c2 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/S3BatchUploader.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/service/provider/S3BatchUploader.java @@ -2,59 +2,64 @@ import io.smallrye.mutiny.Uni; import java.net.URI; +import java.nio.charset.StandardCharsets; import java.time.Duration; +import java.time.Instant; +import java.time.ZoneOffset; +import java.time.format.DateTimeFormatter; +import java.util.List; import java.util.Objects; import java.util.Optional; -import software.amazon.awssdk.auth.credentials.DefaultCredentialsProvider; +import java.util.UUID; import software.amazon.awssdk.core.async.AsyncRequestBody; import software.amazon.awssdk.regions.Region; import software.amazon.awssdk.services.s3.S3AsyncClient; import software.amazon.awssdk.services.s3.model.PutObjectRequest; -/** - * {@link BillingS3LogHandler.AsyncBatchUploader} backed by an AWS SDK v2 {@link S3AsyncClient}: - * each sealed batch is one async {@code PutObject} with bounded retry/backoff, returned as a {@link - * Uni} so the handler's pipeline drives it without blocking. - * - *

Credentials come from the {@link DefaultCredentialsProvider} chain (IRSA web-identity token in - * AWS deployments). When an {@code endpointOverride} is configured (e.g. S3Mock in tests), - * path-style addressing is forced so bucket-as-host resolution does not get in the way. - */ +/** Uploads sealed billing batches to S3 as NDJSON objects under time-partitioned keys. */ public class S3BatchUploader implements BillingS3LogHandler.AsyncBatchUploader { + // S3 object-key consistent identifier; TBD + static final String PATH_PREFIX = "billing-events"; private static final String NDJSON_CONTENT_TYPE = "application/x-ndjson"; + // object key format + private static final DateTimeFormatter KEY_TIME_FORMAT = + DateTimeFormatter.ofPattern("yyyy/MM/dd/HH/mm").withZone(ZoneOffset.UTC); + + // Bound every PUT so a hung connection can neither pin an upload slot indefinitely nor stall + // the shutdown drain. Retries stay inside the SDK's built-in default policy (bounded attempts, + // jittered throttle-aware backoff). + private static final Duration API_CALL_ATTEMPT_TIMEOUT = Duration.ofSeconds(10); + private static final Duration API_CALL_TIMEOUT = Duration.ofSeconds(30); private final S3AsyncClient client; private final String bucket; - private final RetryPolicy retry; - S3BatchUploader(S3AsyncClient client, String bucket, RetryPolicy retry) { + S3BatchUploader(S3AsyncClient client, String bucket) { this.client = client; this.bucket = bucket; - this.retry = retry; } - /** - * Builds an uploader from resolved inputs. {@code endpointOverride} is present only for a non-AWS - * S3 (e.g. S3Mock in tests). - */ public static S3BatchUploader create( - String region, String bucket, Optional endpointOverride, RetryPolicy retry) { + String region, String bucket, Optional endpointOverride) { if (region == null || region.isBlank()) throw new IllegalArgumentException("stargate.jsonapi.billing.s3.bucket-region must be set"); if (bucket == null || bucket.isBlank()) throw new IllegalArgumentException("stargate.jsonapi.billing.s3.bucket must be set"); Objects.requireNonNull(endpointOverride, "endpointOverride must not be null"); - Objects.requireNonNull(retry, "retry must not be null"); + // Credentials resolve from the SDK's default provider chain (env vars, web-identity/OIDC + // token, instance/container roles), left implicit so the client owns — and closes — the + // provider. This transparently supports federated (AssumeRoleWithWebIdentity) and + // cross-account access: the bucket may live in a different account (per IAM + bucket + // policy); its region is set via .region(). var builder = S3AsyncClient.builder() .region(Region.of(region)) - // Credentials resolve from the SDK's default provider chain (env vars, - // web-identity/OIDC token, instance/container roles). This transparently supports - // federated (AssumeRoleWithWebIdentity) and cross-account access — the bucket may live - // in a different account (per IAM + bucket policy); its region is set via .region(). - .credentialsProvider(DefaultCredentialsProvider.create()); + .overrideConfiguration( + o -> + o.apiCallAttemptTimeout(API_CALL_ATTEMPT_TIMEOUT) + .apiCallTimeout(API_CALL_TIMEOUT)); // Real AWS S3 needs no endpoint: the SDK endpoint rules (s3 SDK's DefaultS3EndpointProvider) // derive https://.s3..amazonaws.com from region + partition dnsSuffix. @@ -64,11 +69,13 @@ public static S3BatchUploader create( .filter(s -> !s.isBlank()) .ifPresent(uri -> builder.endpointOverride(URI.create(uri)).forcePathStyle(true)); - return new S3BatchUploader(builder.build(), bucket, retry); + return new S3BatchUploader(builder.build(), bucket); } @Override - public Uni upload(String key, byte[] body) { + public Uni upload(BillingQueue.Batch batch) { + String key = objectKey(batch.firstEventAt(), UUID.randomUUID()); + byte[] body = toNdjson(batch.lines()); return Uni.createFrom() .completionStage( () -> @@ -80,38 +87,23 @@ public Uni upload(String key, byte[] body) { .contentType(NDJSON_CONTENT_TYPE) .build(), AsyncRequestBody.fromBytes(body)) - .thenAccept(resp -> {})) - .onFailure() - .retry() - .withBackOff(retry.initialBackOff(), retry.maxBackOff()) - .withJitter(retry.jitter()) - .atMost(retry.atMostRetries()); + .thenAccept(resp -> {})); } - @Override - public void close() { - client.close(); + static String objectKey(Instant timestamp, UUID id) { + return PATH_PREFIX + "/" + KEY_TIME_FORMAT.format(timestamp) + "/" + id + ".jsonl"; } - /** - * Bounded exponential-backoff-with-jitter tuning for one PUT's retries. Validated on construction - * so bad tuning fails at wiring time with a clear message. - */ - public record RetryPolicy( - int atMostRetries, Duration initialBackOff, Duration maxBackOff, double jitter) { - public RetryPolicy { - if (atMostRetries < 0) { - throw new IllegalArgumentException("atMostRetries must be >= 0"); - } - if (initialBackOff.isNegative() || initialBackOff.isZero()) { - throw new IllegalArgumentException("initialBackOff must be > 0"); - } - if (maxBackOff.compareTo(initialBackOff) < 0) { - throw new IllegalArgumentException("maxBackOff must be >= initialBackOff"); - } - if (jitter < 0 || jitter > 1) { - throw new IllegalArgumentException("jitter must be in [0, 1]"); - } + static byte[] toNdjson(List lines) { + StringBuilder sb = new StringBuilder(); + for (String line : lines) { + sb.append(line).append('\n'); } + return sb.toString().getBytes(StandardCharsets.UTF_8); + } + + @Override + public void close() { + client.close(); } } From 87baefd4aa30b4ed09b092eed0f26ca387117dda Mon Sep 17 00:00:00 2001 From: Hazel Date: Tue, 14 Jul 2026 19:42:08 -0700 Subject: [PATCH 20/65] Update BillingQueue --- .../service/provider/BillingQueue.java | 107 ++++++++++++++++++ 1 file changed, 107 insertions(+) create mode 100644 src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingQueue.java diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingQueue.java b/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingQueue.java new file mode 100644 index 0000000000..b11dd6401d --- /dev/null +++ b/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingQueue.java @@ -0,0 +1,107 @@ +package io.stargate.sgv2.jsonapi.service.provider; + +import java.time.Instant; +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.ArrayBlockingQueue; +import java.util.concurrent.BlockingQueue; +import java.util.concurrent.atomic.AtomicLong; + +/** + * Bounded in-memory buffer that owns the batching policy of the billing S3 export: it decides when + * a batch is sealed ({@code maxEvents} line-count seal, {@code maxBytes} body-size seal) and hands + * out drained {@link Batch}es. + */ +public final class BillingQueue { + + private final BlockingQueue queue; + // Approximate buffered NDJSON bytes (see lineBytes); + private final AtomicLong queuedBytes = new AtomicLong(0); + private final int batchSize; + private final long maxBytes; + + public BillingQueue(int maxEvents, long maxBytes, int queueCapacity) { + if (maxEvents < 1) { + throw new IllegalArgumentException( + "stargate.jsonapi.billing.s3.max-events must be >= 1 (was " + maxEvents + ")"); + } + if (maxBytes < 1) { + throw new IllegalArgumentException( + "stargate.jsonapi.billing.s3.max-bytes must be >= 1 (was " + maxBytes + ")"); + } + if (queueCapacity < 1) { + throw new IllegalArgumentException( + "stargate.jsonapi.billing.s3.queue-capacity must be >= 1 (was " + queueCapacity + ")"); + } + this.batchSize = maxEvents; + this.maxBytes = maxBytes; + this.queue = new ArrayBlockingQueue<>(queueCapacity); + } + + /** + * Buffers one line, or returns {@code false} when the capacity bound is hit. + * + * @param eventAt when the event was logged; carried through to {@link Batch#oldestEventAt()} + */ + public boolean offer(Instant eventAt, String line) { + if (!queue.offer(new Entry(eventAt, line))) { + return false; + } + queuedBytes.addAndGet(lineBytes(line)); + return true; + } + + /** True once a seal is reached: a full batch by line count, or {@code maxBytes} buffered. */ + public boolean shouldFlush() { + return queue.size() >= batchSize || queuedBytes.get() >= maxBytes; + } + + public boolean isEmpty() { + return queue.isEmpty(); + } + + public int size() { + return queue.size(); + } + + /** Removes and returns up to one sealed batch (possibly partial, possibly {@code EMPTY}). */ + public Batch drain() { + List lines = new ArrayList<>(batchSize); + Instant oldestEventAt = null; + long bytes = 0; + Entry entry; + while (lines.size() < batchSize && bytes < maxBytes && (entry = queue.poll()) != null) { + if (oldestEventAt == null || entry.eventAt().isBefore(oldestEventAt)) { + oldestEventAt = entry.eventAt(); + } + lines.add(entry.line()); + bytes += lineBytes(entry.line()); + } + queuedBytes.addAndGet(-bytes); + return oldestEventAt == null ? Batch.EMPTY : new Batch(lines, oldestEventAt); + } + + // String.length() (UTF-16 units) + newline as a cheap stand-in for UTF-8 bytes: exact for the + // ASCII JSON , an undercount for non-ASCII + private static int lineBytes(String line) { + return line.length() + 1; + } + + /** + * One drained, sealed batch. {@code oldestEventAt} is the minimum event time across {@code lines} + * — queue order is enqueue order, not event-time order, under concurrent publish. + */ + public record Batch(List lines, Instant oldestEventAt) { + static final Batch EMPTY = new Batch(List.of(), Instant.EPOCH); + + boolean isEmpty() { + return lines.isEmpty(); + } + + int events() { + return lines.size(); + } + } + + private record Entry(Instant eventAt, String line) {} +} From b8c168720ecd78ba6d3afdc859de4509917e0d0f Mon Sep 17 00:00:00 2001 From: Hazel Date: Tue, 14 Jul 2026 19:42:47 -0700 Subject: [PATCH 21/65] Update batch oldestEventAt --- .../stargate/sgv2/jsonapi/service/provider/S3BatchUploader.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/S3BatchUploader.java b/src/main/java/io/stargate/sgv2/jsonapi/service/provider/S3BatchUploader.java index 31532bb7c2..2043a7bc5b 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/S3BatchUploader.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/service/provider/S3BatchUploader.java @@ -74,7 +74,7 @@ public static S3BatchUploader create( @Override public Uni upload(BillingQueue.Batch batch) { - String key = objectKey(batch.firstEventAt(), UUID.randomUUID()); + String key = objectKey(batch.oldestEventAt(), UUID.randomUUID()); byte[] body = toNdjson(batch.lines()); return Uni.createFrom() .completionStage( From bfe18148c4059cb7e756cf5f465efa44b266fc62 Mon Sep 17 00:00:00 2001 From: Hazel Date: Tue, 14 Jul 2026 19:44:37 -0700 Subject: [PATCH 22/65] Update uploader --- .../service/provider/BillingS3HandlerInstaller.java | 12 +----------- 1 file changed, 1 insertion(+), 11 deletions(-) diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3HandlerInstaller.java b/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3HandlerInstaller.java index b912a45809..c1057551e3 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3HandlerInstaller.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3HandlerInstaller.java @@ -7,7 +7,6 @@ import jakarta.enterprise.context.ApplicationScoped; import jakarta.enterprise.event.Observes; import jakarta.inject.Inject; -import java.time.Duration; import java.util.logging.Logger; import org.slf4j.LoggerFactory; @@ -53,16 +52,7 @@ void onStart(@Observes StartupEvent event) { var bucket = config.bucket().orElse(null); // Fail-loud: invalid billing S3 config throws here, aborting application startup. - var uploader = - S3BatchUploader.create( - region, - bucket, - config.endpointOverride(), - new S3BatchUploader.RetryPolicy( - config.atMostRetries(), - Duration.ofMillis(config.initialBackOffMillis()), - Duration.ofMillis(config.maxBackOffMillis()), - config.retryJitter())); + var uploader = S3BatchUploader.create(region, bucket, config.endpointOverride()); this.handler = new BillingS3LogHandler(config, uploader, meterRegistry); Logger.getLogger(BILLING_LOGGER_NAME).addHandler(this.handler); From 720e34a376ed0878ffaf0ba2faaf07f5e6a84b18 Mon Sep 17 00:00:00 2001 From: Hazel Date: Tue, 14 Jul 2026 20:24:54 -0700 Subject: [PATCH 23/65] Update BillingS3LogHandler --- .../service/provider/BillingS3LogHandler.java | 287 +++++++----------- 1 file changed, 109 insertions(+), 178 deletions(-) diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandler.java b/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandler.java index 2678633d5d..1780d31e19 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandler.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandler.java @@ -1,25 +1,14 @@ package io.stargate.sgv2.jsonapi.service.provider; -import com.fasterxml.jackson.databind.JsonNode; -import com.fasterxml.jackson.databind.ObjectMapper; import com.google.common.annotations.VisibleForTesting; import io.micrometer.core.instrument.MeterRegistry; import io.smallrye.mutiny.Uni; import io.smallrye.mutiny.infrastructure.Infrastructure; import io.stargate.sgv2.jsonapi.config.BillingS3ExportConfig; -import java.nio.charset.StandardCharsets; import java.time.Duration; -import java.time.Instant; -import java.time.ZoneOffset; -import java.time.format.DateTimeFormatter; -import java.util.ArrayList; -import java.util.List; -import java.util.UUID; -import java.util.concurrent.ArrayBlockingQueue; -import java.util.concurrent.BlockingQueue; +import java.util.concurrent.ScheduledFuture; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicInteger; -import java.util.concurrent.atomic.AtomicLong; import java.util.concurrent.locks.LockSupport; import java.util.logging.Handler; import java.util.logging.LogRecord; @@ -27,64 +16,29 @@ import org.slf4j.LoggerFactory; /** - * A {@link Handler} that ships {@code billing.events} JSON log lines to S3 as NDJSON ({@code - * .jsonl}) objects. Installed on the {@code billing.events} logger by {@link - * BillingS3HandlerInstaller} when {@link BillingS3ExportConfig#enabled()} is {@code true}; the - * existing console handler stays attached as a backstop (dual-write). + * Buffers {@code billing.events} log lines and ships them to S3 in sealed batches. * - *

Off the request path. {@link #publish(LogRecord)} only offers the line to a bounded - * in-memory queue — it never blocks and never throws. When the queue holds a full batch ({@link - * BillingS3ExportConfig#maxEvents()} lines) a flush is dispatched to a worker-pool thread: it - * drains up to that many lines and PUTs them as one object. Up to {@link - * BillingS3ExportConfig#uploadConcurrency()} flushes run at once — a non-blocking in-flight counter - * is the gate, and the S3 PUT is async so worker threads are not held during upload. When the queue - * is full, further lines are dropped and counted (never silent). - * - *

Verbatim bodies. Each log line is kept byte-for-byte as one NDJSON row; only the - * batch's first line is parsed (for the key's date path). Each flushed batch is one object at - * {@code ///

///.jsonl}; the key is built once and reused across - * the uploader's retries so a retried PUT overwrites rather than duplicates. - * - *

{@link #close()} drains whatever remains (including a final under-batch tail), waits for - * in-flight flushes to settle, then closes the uploader. The handler is intentionally not a - * CDI bean — the installer wires this instance to the {@code billing.events} category explicitly. + *

Pure orchestration: {@link BillingQueue} owns the batching policy (when a batch seals), the + * {@link AsyncBatchUploader} owns what lands in the bucket (key layout, body encoding), and this + * handler wires them together — a non-blocking publish path, seal- and age-triggered flushes under + * bounded upload concurrency, metrics, and a deadline-bounded drain on close. */ public final class BillingS3LogHandler extends Handler { - // ---- Constants ---- - // S3 object-key consistent identifier; TBD - static final String PATH_PREFIX = "billing-events"; private static final Logger LOG = LoggerFactory.getLogger(BillingS3LogHandler.class); - private static final ObjectMapper MAPPER = new ObjectMapper(); - // UTC, minute-resolution date path for the object key - private static final DateTimeFormatter KEY_TIME_FORMAT = - DateTimeFormatter.ofPattern("yyyy/MM/dd/HH/mm").withZone(ZoneOffset.UTC); // ---- Collaborators ---- private final AsyncBatchUploader uploader; private final BillingMetrics metrics; + private final BillingQueue buffer; - // ---- Tuning (resolved from BillingS3ExportConfig) ---- - private final int batchSize; // flush trigger + max lines per object (config.maxEvents) - private final long maxBytes; // max NDJSON body bytes per object (config.maxBytes) + // ---- Flush pipeline ---- private final int uploadConcurrency; // max flushes (S3 PUTs) in flight at once - - // ---- Queue + byte/concurrency gates ---- - private final BlockingQueue queue; - - /** - * Running total of NDJSON body bytes buffered in {@link #queue} (each line's length plus its - * newline) — drives the byte-based flush trigger. - */ - private final AtomicLong queuedBytes = new AtomicLong(0); - - /** - * Flushes (S3 PUTs) currently in flight — the non-blocking concurrency gate, capped at {@link - * #uploadConcurrency}. A slot is held from {@link #drain} through the PUT's completion. - */ private final AtomicInteger inFlight = new AtomicInteger(0); + private final ScheduledFuture ageFlushTask; + + private final Duration shutdownTimeout; - /** Config-driven constructor used by the installer. */ public BillingS3LogHandler( BillingS3ExportConfig config, AsyncBatchUploader uploader, MeterRegistry meterRegistry) { this( @@ -94,15 +48,10 @@ public BillingS3LogHandler( config.maxBytes(), config.maxAge(), config.queueCapacity(), - config.uploadConcurrency()); + config.uploadConcurrency(), + config.shutdownTimeout()); } - /** - * Explicit-threshold constructor; convenient for unit tests. {@code maxAge} is accepted for - * config compatibility but not yet wired (time-based flush is a follow-up); a batch is sealed on - * whichever of {@code maxEvents} (line count) or {@code maxBytes} (NDJSON body size) it hits - * first. - */ @VisibleForTesting BillingS3LogHandler( AsyncBatchUploader uploader, @@ -111,13 +60,30 @@ public BillingS3LogHandler( long maxBytes, Duration maxAge, int queueCapacity, - int uploadConcurrency) { + int uploadConcurrency, + Duration shutdownTimeout) { + if (uploadConcurrency < 1) { + throw new IllegalArgumentException( + "s3.upload-concurrency must be >= 1 (was " + uploadConcurrency + ")"); + } + requirePositive(maxAge, "max-age"); + requirePositive(shutdownTimeout, "shutdown-timeout"); this.uploader = uploader; - this.batchSize = maxEvents; - this.maxBytes = maxBytes; this.uploadConcurrency = uploadConcurrency; - this.queue = new ArrayBlockingQueue<>(queueCapacity); - this.metrics = new BillingMetrics(meterRegistry, queue::size, queueCapacity); + this.shutdownTimeout = shutdownTimeout; + this.buffer = new BillingQueue(maxEvents, maxBytes, queueCapacity); + this.metrics = new BillingMetrics(meterRegistry, buffer::size, queueCapacity); + this.ageFlushTask = + Infrastructure.getDefaultWorkerPool() + .scheduleAtFixedRate( + this::onAgeTick, maxAge.toMillis(), maxAge.toMillis(), TimeUnit.MILLISECONDS); + } + + private static void requirePositive(Duration value, String property) { + if (value == null || value.isNegative() || value.isZero()) { + throw new IllegalArgumentException( + "stargate.jsonapi.billing.s3." + property + " must be > 0 (was " + value + ")"); + } } // ============================================================ @@ -129,173 +95,138 @@ public void publish(LogRecord record) { if (record == null) { return; } + // Producer contract (DefaultBilling): the message is the final JSON line, logged without {} + // placeholders. This handler never runs a Formatter, so a parameterized call would ship its + // raw template. getInstant() is when the producer logged it — within microseconds of the + // "timestamp" it embedded in the JSON, and the object key only needs minute resolution. String line = record.getMessage(); if (line == null || line.isBlank()) { return; } metrics.recordOffered(); - if (!queue.offer(line)) { - // Bounded queue full: shed and count — never block, never throw, never silent. + if (!buffer.offer(record.getInstant(), line)) { + // Bounded buffer full: drop and count metrics.recordDropped(); return; } - queuedBytes.addAndGet(lineBytes(line)); maybeFlush(); } @Override public void flush() { - // No-op: flushes are size-triggered and continuous; close() seals whatever remains on shutdown. + // No-op: shipping is seal-triggered (publish) and age-triggered (tick); close() drains. } /** - * Shutdown drain: serially flushes the buffered backlog (incl. the sub-batch tail the size-gate - * skips), then waits for in-flight PUTs to finish before closing the client. No internal timeout - * — bounded by the platform's shutdown grace (SIGKILL). + * Drains what remains through the normal flush pipeline, bounded by {@code shutdownTimeout}. The + * budget only bites when S3 is already failing: it converts a silent SIGKILL into a logged count + * of abandoned events and lets the rest of shutdown proceed. */ @Override public void close() { - List batch; - while (!(batch = drain()).isEmpty()) { - uploadBatch(batch).await().indefinitely(); + // Don't interrupt a tick already running; the in-flight wait below covers it. + ageFlushTask.cancel(false); + long deadlineNanos = System.nanoTime() + shutdownTimeout.toNanos(); + + // Pump the pipeline until the buffer is drained: tryFlush() is a no-op while all slots are + // busy, and every settled upload frees a slot for the next batch. + while (!buffer.isEmpty() && System.nanoTime() < deadlineNanos) { + tryFlush(); + LockSupport.parkNanos(TimeUnit.MILLISECONDS.toNanos(1)); } - while (inFlight.get() > 0) { + // Let in-flight uploads (ours and any started before close) settle within the budget. + while (inFlight.get() > 0 && System.nanoTime() < deadlineNanos) { LockSupport.parkNanos(TimeUnit.MILLISECONDS.toNanos(10)); } - uploader.close(); + + int queuedAbandoned = buffer.size(); + int inFlightAbandoned = inFlight.get(); + if (queuedAbandoned > 0 || inFlightAbandoned > 0) { + LOG.warn( + "Billing S3 export shutdown budget ({}) exhausted: dropping {} buffered events," + + " abandoning {} in-flight uploads", + shutdownTimeout, + queuedAbandoned, + inFlightAbandoned); + } + uploader.close(); // aborts anything still in flight } // ============================================================ // Flush pipeline // ============================================================ - /** - * Dispatches one flush if a full object's worth is buffered ({@link #shouldFlush}) and a - * concurrency slot is free. Called after every {@link #publish} and again when a flush completes, - * so the pipeline self-clocks up to {@link #uploadConcurrency} concurrent uploads with no - * standing reader thread. - */ + /** Seal-triggered flush: ship when the buffer has a full batch by count or bytes. */ private void maybeFlush() { - if (shouldFlush() && tryFlush()) { - Uni.createFrom() - .item(this::drain) - .runSubscriptionOn(Infrastructure.getDefaultWorkerPool()) - .flatMap(this::uploadBatch) - .eventually( - () -> { - inFlight.getAndDecrement(); - maybeFlush(); // a full batch may have accumulated while this one uploaded - }) - .subscribe() - .with(ignored -> {}, failure -> LOG.error("Billing S3 export flush failed", failure)); + if (buffer.shouldFlush()) { + tryFlush(); } } - /** True when the queue holds a full object's worth — by line count or by buffered bytes. */ - private boolean shouldFlush() { - return queue.size() >= batchSize || queuedBytes.get() >= maxBytes; - } - /** - * Non-blocking CAS gate: claims an in-flight slot iff fewer than {@link #uploadConcurrency} are - * held. + * Age-triggered flush: every {@code maxAge} tick ships whatever is buffered, sealed or not, so no + * event waits longer than ~{@code maxAge} even under trickle traffic. */ - private boolean tryFlush() { - int prev = inFlight.getAndUpdate(n -> n < uploadConcurrency ? n + 1 : n); - return prev < uploadConcurrency; - } - - /** - * Removes lines from the queue for one object — up to {@link #batchSize} lines or {@link - * #maxBytes} of NDJSON body, whichever comes first (a single over-cap line still goes out alone). - * May come back empty if another flush raced ahead and took them first. - */ - private List drain() { - List batch = new ArrayList<>(batchSize); - long bytes = 0; - String line; - while (batch.size() < batchSize && bytes < maxBytes && (line = queue.poll()) != null) { - batch.add(line); - bytes += lineBytes(line); + private void onAgeTick() { + try { + if (!buffer.isEmpty()) { + tryFlush(); + } + } catch (Throwable t) { + LOG.error("Billing S3 export age-flush tick failed", t); } - queuedBytes.addAndGet(-bytes); - return batch; } /** - * NDJSON body bytes one line contributes: its length (== UTF-8 bytes for ASCII) plus a newline. + * Claims an in-flight slot (non-blocking CAS, at most {@link #uploadConcurrency} held) and, on + * success, drains + uploads one batch asynchronously. When the upload settles the slot is + * released and the seal condition re-checked: a full batch may have accumulated meanwhile. */ - private static int lineBytes(String line) { - return line.length() + 1; + private void tryFlush() { + int prev = inFlight.getAndUpdate(n -> n < uploadConcurrency ? n + 1 : n); + if (prev >= uploadConcurrency) { + return; + } + Uni.createFrom() + .item(buffer::drain) + .runSubscriptionOn(Infrastructure.getDefaultWorkerPool()) + .flatMap(this::uploadBatch) + .eventually( + () -> { + inFlight.getAndDecrement(); + maybeFlush(); + }) + .subscribe() + .with(ignored -> {}, failure -> LOG.error("Billing S3 export flush failed", failure)); } - /** - * Ships one drained batch: the uploader PUTs it as a single NDJSON object, applying its own - * bounded retry/backoff. Never propagates failure — a giving-up batch is counted and recovered to - * a no-op so the flush loop stays alive. - */ - private Uni uploadBatch(List batch) { + /** Uploads one batch; never fails the pipeline — a batch that exhausts retries is counted. */ + private Uni uploadBatch(BillingQueue.Batch batch) { if (batch.isEmpty()) { return Uni.createFrom().voidItem(); } - int events = batch.size(); - String key = objectKey(firstTimestamp(batch), UUID.randomUUID()); - byte[] body = toNdjson(batch); + int size = batch.size(); return uploader - .upload(key, body) + .upload(batch) .onItem() - .invoke(() -> metrics.recordBatchDelivered(events)) + .invoke(() -> metrics.recordBatchDelivered(size)) .onFailure() - .invoke( - t -> LOG.error("Failed to upload billing S3 batch '{}' ({} events)", key, events, t)) + .invoke(t -> LOG.error("Failed to upload billing S3 batch ({} size)", size, t)) .onFailure() .recoverWithItem( () -> { - metrics.recordBatchFailed(events); + metrics.recordBatchFailed(size); return null; }); } - /** Object key: {@code ///

///.jsonl} (UTC). */ - static String objectKey(Instant timestamp, UUID id) { - return PATH_PREFIX + "/" + KEY_TIME_FORMAT.format(timestamp) + "/" + id + ".jsonl"; - } - - /** NDJSON body: each line verbatim, newline-terminated. */ - private static byte[] toNdjson(List batch) { - StringBuilder sb = new StringBuilder(); - for (String line : batch) { - sb.append(line).append('\n'); - } - return sb.toString().getBytes(StandardCharsets.UTF_8); - } - - /** - * Timestamp for the object key's date path — parsed from the batch's first line's {@code - * timestamp}, falling back to wall clock (and counting the parse failure) when it is - * absent/unparseable. - */ - private Instant firstTimestamp(List batch) { - try { - JsonNode node = MAPPER.readTree(batch.get(0)); - JsonNode tsNode = node.get("timestamp"); - if (tsNode != null && tsNode.isTextual()) { - return Instant.parse(tsNode.asText()); - } - } catch (Exception e) { - // fall through to the wall-clock fallback below - } - metrics.recordParseFailure(); - return Instant.now(); - } - /** - * Uploads one drained batch to S3 with its own bounded retry/backoff, failing only once retries - * are exhausted (see implementation {@link S3BatchUploader}). + * Uploads one sealed batch to the export destination; owns the object key and body encoding. + * Implementations must tolerate concurrent calls. */ @FunctionalInterface public interface AsyncBatchUploader extends AutoCloseable { - Uni upload(String key, byte[] body); + Uni upload(BillingQueue.Batch batch); @Override default void close() {} From cd896868961d58607f8970bcc763a3c277cfcfaa Mon Sep 17 00:00:00 2001 From: Hazel Date: Tue, 14 Jul 2026 20:31:27 -0700 Subject: [PATCH 24/65] Update size() --- .../io/stargate/sgv2/jsonapi/service/provider/BillingQueue.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingQueue.java b/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingQueue.java index b11dd6401d..873e0df20d 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingQueue.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingQueue.java @@ -98,7 +98,7 @@ boolean isEmpty() { return lines.isEmpty(); } - int events() { + int size() { return lines.size(); } } From eb191428431f4379f92192804d6888391f65e472 Mon Sep 17 00:00:00 2001 From: Hazel Date: Tue, 14 Jul 2026 20:55:33 -0700 Subject: [PATCH 25/65] Update BillingS3ExportConfig --- .../jsonapi/config/BillingS3ExportConfig.java | 70 ++++--------------- 1 file changed, 14 insertions(+), 56 deletions(-) diff --git a/src/main/java/io/stargate/sgv2/jsonapi/config/BillingS3ExportConfig.java b/src/main/java/io/stargate/sgv2/jsonapi/config/BillingS3ExportConfig.java index f1f910dd66..f8b927477f 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/config/BillingS3ExportConfig.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/config/BillingS3ExportConfig.java @@ -5,86 +5,44 @@ import java.time.Duration; import java.util.Optional; -/** - * Configuration for exporting {@code billing.events} log lines to S3 as NDJSON {@code .jsonl} - * objects. Consumed by {@link io.stargate.sgv2.jsonapi.service.provider.BillingS3HandlerInstaller} - * which, when {@link #enabled()} is {@code true}, attaches a {@link - * io.stargate.sgv2.jsonapi.service.provider.BillingS3LogHandler} to the {@code billing.events} - * logger. - * - *

This is a startup-time switch, not a per-request feature flag — it is independent of - * {@link io.stargate.sgv2.jsonapi.config.feature.ApiFeature#BILLING_EVENTS_LOGGING}. Events only - * reach the handler when the {@code billing.events} logger is also emitting (i.e. the billing - * feature is on); this flag then decides whether those lines are additionally shipped to S3. The - * existing console handler stays attached as a backstop regardless. - * - *

Off by default. When enabled, {@link #bucket()} and {@link #bucketRegion()} are - * required; if either is missing the handler is not installed (logged as an error) and billing - * events continue to flow to the console only. - */ +/** Configuration for the billing S3 export (see BillingS3HandlerInstaller). */ @ConfigMapping(prefix = "stargate.jsonapi.billing.s3") public interface BillingS3ExportConfig { - /** - * Master switch; when {@code false} (default) no handler is installed and no S3 client is built. - */ + /** Master switch: when false the export handler is never installed. */ @WithDefault("false") boolean enabled(); - /** Target bucket, e.g. {@code serverless-usage-dev}. Required when {@link #enabled()}. */ + /** S3 bucket name */ Optional bucket(); - /** AWS region of the bucket, e.g. {@code us-east-1}. Required when {@link #enabled()}. */ + /** S3 bucket region */ Optional bucketRegion(); - /** - * Endpoint override for the S3 client. Set this to point at a non-AWS S3 (e.g. S3Mock in tests); - * SDK resolves the regional AWS endpoint when left empty. - */ + /** Only for non-AWS S3 endpoints (e.g. S3Mock in tests). */ Optional endpointOverride(); - /** Seal a batch once it holds this many events. */ + /** Line-count seal: a buffered batch is shipped once it holds this many events. */ @WithDefault("50") int maxEvents(); - /** Seal a batch once its NDJSON body reaches this many bytes (~2 MiB default). */ + /** Byte-size seal on the buffered NDJSON body. */ @WithDefault("2097152") long maxBytes(); - /** Seal an open (under-filled) batch once its oldest event is this old (flush interval). */ + /** Age flush period: buffered events are shipped at least this often, sealed or not. */ @WithDefault("PT30S") Duration maxAge(); - /** - * Capacity of the handler's in-memory hand-off queue. {@link BillingS3LogHandler#publish()} - * offers lines non-blocking; once the queue is full, further lines are dropped and counted. - */ + /** Bound on buffered events; beyond it new lines are dropped. */ @WithDefault("10000") int queueCapacity(); - /** - * Maximum number of retries after a failed PUT, per sealed batch, before the batch is counted as - * failed. Default is 2 retries (up to 3 attempts including the initial PUT). - */ - @WithDefault("2") - int atMostRetries(); - - /** - * The initial delay between retries in milliseconds. The first retry occurs after the specified - * delay (default 100 ms), doubling each time until reaching maxBackOffMillis. - */ - @WithDefault("100") - int initialBackOffMillis(); - - /** The maximum delay between retries in milliseconds. */ - @WithDefault("500") - int maxBackOffMillis(); - - /** A random variation added to the delay between retries in an exponential backoff strategy. */ - @WithDefault("0.5") - double retryJitter(); - - /** Number of batch uploads (S3 PUTs) allowed in flight concurrently. */ + /** Max concurrent S3 PUTs. */ @WithDefault("4") int uploadConcurrency(); + + /** Budget for draining the buffer at shutdown; keep below the pod termination grace period. */ + @WithDefault("PT20S") + Duration shutdownTimeout(); } From d2f06c2ded06a682bc6ee0c6efe197b2d982fadc Mon Sep 17 00:00:00 2001 From: Hazel Date: Tue, 14 Jul 2026 21:34:16 -0700 Subject: [PATCH 26/65] Update BillingS3LogHandler --- .../service/provider/BillingS3LogHandler.java | 28 +++++++++++++------ 1 file changed, 20 insertions(+), 8 deletions(-) diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandler.java b/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandler.java index 1780d31e19..a7e40abd10 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandler.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandler.java @@ -16,12 +16,15 @@ import org.slf4j.LoggerFactory; /** - * Buffers {@code billing.events} log lines and ships them to S3 in sealed batches. + * JUL handler that turns {@code billing.events} log lines into batched S3 objects. * - *

Pure orchestration: {@link BillingQueue} owns the batching policy (when a batch seals), the - * {@link AsyncBatchUploader} owns what lands in the bucket (key layout, body encoding), and this - * handler wires them together — a non-blocking publish path, seal- and age-triggered flushes under - * bounded upload concurrency, metrics, and a deadline-bounded drain on close. + *

Division of labor: {@link BillingQueue} decides when a batch seals, {@link AsyncBatchUploader} + * decides what an S3 object looks like, and this class decides when uploads run — the flush + * triggers (seal on publish, age tick, drain on close), the upload-concurrency gate, and metrics. + * + *

Delivery is at-most-once by design: publish never blocks and never throws — when the buffer is + * full, lines are dropped and counted — and close() gives up loudly once {@code shutdownTimeout} is + * exhausted. */ public final class BillingS3LogHandler extends Handler { @@ -64,7 +67,9 @@ public BillingS3LogHandler( Duration shutdownTimeout) { if (uploadConcurrency < 1) { throw new IllegalArgumentException( - "s3.upload-concurrency must be >= 1 (was " + uploadConcurrency + ")"); + "stargate.jsonapi.billing.s3.upload-concurrency must be >= 1 (was " + + uploadConcurrency + + ")"); } requirePositive(maxAge, "max-age"); requirePositive(shutdownTimeout, "shutdown-timeout"); @@ -142,6 +147,8 @@ public void close() { int queuedAbandoned = buffer.size(); int inFlightAbandoned = inFlight.get(); if (queuedAbandoned > 0 || inFlightAbandoned > 0) { + // Only queued events are counted: in-flight ones settle as failed when the abort lands. + metrics.recordAbandonedAtShutdown(queuedAbandoned); LOG.warn( "Billing S3 export shutdown budget ({}) exhausted: dropping {} buffered events," + " abandoning {} in-flight uploads", @@ -164,8 +171,13 @@ private void maybeFlush() { } /** - * Age-triggered flush: every {@code maxAge} tick ships whatever is buffered, sealed or not, so no - * event waits longer than ~{@code maxAge} even under trickle traffic. + * Age trigger: every {@code maxAge} tick ships whatever is buffered, sealed or not. Deliberately + * no head-age check: flushing only entries older than {@code maxAge} would let an event that just + * missed a tick wait ~2x{@code maxAge}, while shipping unconditionally bounds every wait by one + * period — at the cost of an occasional small object when a tick lands just after a seal flush. + * + *

Catches everything: an escaped throwable would silently cancel all future runs of a + * fixed-rate task. */ private void onAgeTick() { try { From ed283469af71c9b77a070aa6389697a1b207da4c Mon Sep 17 00:00:00 2001 From: Hazel Date: Tue, 14 Jul 2026 22:07:54 -0700 Subject: [PATCH 27/65] Add BillingMetrics --- .../service/provider/BillingMetrics.java | 98 +++++++++++++++++++ 1 file changed, 98 insertions(+) create mode 100644 src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingMetrics.java diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingMetrics.java b/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingMetrics.java new file mode 100644 index 0000000000..b2198b550c --- /dev/null +++ b/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingMetrics.java @@ -0,0 +1,98 @@ +package io.stargate.sgv2.jsonapi.service.provider; + +import io.micrometer.core.instrument.Counter; +import io.micrometer.core.instrument.Gauge; +import io.micrometer.core.instrument.MeterRegistry; +import java.time.Instant; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicLong; +import java.util.function.Supplier; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Meters for the billing S3 export. Every offered event ends in exactly one terminal counter — + * {@code flushed}, {@code failed}, or {@code dropped(reason=capacity|shutdown)} — or is still + * visible in {@code queue.depth} / an in-flight batch, so loss is always accounted for. {@code + * last_delivery.epoch_seconds} is the freshness heartbeat to alert on. + */ +public final class BillingMetrics { + + private static final Logger LOG = LoggerFactory.getLogger(BillingMetrics.class); + private static final long DROP_WARN_INTERVAL_NANOS = TimeUnit.MINUTES.toNanos(10); + + private final Counter offered; + private final Counter droppedCapacity; + private final Counter droppedShutdown; + private final Counter flushed; + private final Counter failed; + private final Counter batchesUploaded; + private final Counter batchesFailed; + private final AtomicLong lastDeliveryEpochSeconds = new AtomicLong(0); + private final AtomicLong lastDropWarnNanos; + private final long queueCapacity; + + /** + * @param depthSource live queue depth, exposed read-only as {@code billing.s3.queue.depth} + * @param queueCapacity quoted in the buffer-full warning + */ + public BillingMetrics(MeterRegistry meterRegistry, Supplier depthSource, long queueCapacity) { + this.queueCapacity = queueCapacity; + this.lastDropWarnNanos = new AtomicLong(System.nanoTime() - DROP_WARN_INTERVAL_NANOS); + this.offered = meterRegistry.counter("billing.s3.events.offered"); + this.droppedCapacity = meterRegistry.counter("billing.s3.events.dropped", "reason", "capacity"); + this.droppedShutdown = meterRegistry.counter("billing.s3.events.dropped", "reason", "shutdown"); + this.flushed = meterRegistry.counter("billing.s3.events.flushed"); + this.failed = meterRegistry.counter("billing.s3.events.failed"); + this.batchesUploaded = meterRegistry.counter("billing.s3.batches.uploaded"); + this.batchesFailed = meterRegistry.counter("billing.s3.batches.failed"); + // Backs the {@code last_delivery} gauge; alert on staleness — a dead export and no traffic look identical in counter rates. + Gauge.builder( + "billing.s3.last_delivery.epoch_seconds", + lastDeliveryEpochSeconds, + AtomicLong::doubleValue) + .description("Epoch seconds of the last successful batch delivery") + .register(meterRegistry); + Gauge.builder("billing.s3.queue.depth", depthSource) + .description("Billing events buffered in memory, not yet drained") + .register(meterRegistry); + } + + /** A line was handed to the handler (counted before the capacity check). */ + public void recordOffered() { + offered.increment(); + } + + /** A line was dropped on a full buffer; warns rate-limited so a sustained stall stays visible. */ + public void recordDropped() { + droppedCapacity.increment(); + long now = System.nanoTime(); + long prev = lastDropWarnNanos.get(); + // Rate limit for the buffer-full warning: keeps a drop storm from spamming the logs while + // still surfacing a second stall long after the first. + if (now - prev >= DROP_WARN_INTERVAL_NANOS && lastDropWarnNanos.compareAndSet(prev, now)) { + LOG.warn( + "Billing S3 export backlog full ({} events): shedding billing events because S3 uploads" + + " are slower than ingest. Every shed line is counted by billing.s3.events.dropped.", + queueCapacity); + } + } + + /** Events still buffered when the shutdown budget ran out; close() logs the tombstone. */ + public void recordAbandonedAtShutdown(int size) { + droppedShutdown.increment(size); + } + + /** A batch of events lines landed in S3; bumps the delivery heartbeat. */ + public void recordBatchDelivered(int size) { + flushed.increment(size); + batchesUploaded.increment(); + lastDeliveryEpochSeconds.set(Instant.now().getEpochSecond()); + } + + /** A batch of events lines was given up after the uploader exhausted its retries. */ + public void recordBatchFailed(int size) { + failed.increment(size); + batchesFailed.increment(); + } +} From 8e524d0518dd571cc78b0740042f06b98e83dd83 Mon Sep 17 00:00:00 2001 From: Hazel Date: Tue, 14 Jul 2026 22:27:08 -0700 Subject: [PATCH 28/65] Add BillingMetrics doc --- .../sgv2/jsonapi/service/provider/BillingMetrics.java | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingMetrics.java b/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingMetrics.java index b2198b550c..88d30636ae 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingMetrics.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingMetrics.java @@ -36,7 +36,8 @@ public final class BillingMetrics { * @param depthSource live queue depth, exposed read-only as {@code billing.s3.queue.depth} * @param queueCapacity quoted in the buffer-full warning */ - public BillingMetrics(MeterRegistry meterRegistry, Supplier depthSource, long queueCapacity) { + public BillingMetrics( + MeterRegistry meterRegistry, Supplier depthSource, long queueCapacity) { this.queueCapacity = queueCapacity; this.lastDropWarnNanos = new AtomicLong(System.nanoTime() - DROP_WARN_INTERVAL_NANOS); this.offered = meterRegistry.counter("billing.s3.events.offered"); @@ -46,7 +47,8 @@ public BillingMetrics(MeterRegistry meterRegistry, Supplier depthSource, this.failed = meterRegistry.counter("billing.s3.events.failed"); this.batchesUploaded = meterRegistry.counter("billing.s3.batches.uploaded"); this.batchesFailed = meterRegistry.counter("billing.s3.batches.failed"); - // Backs the {@code last_delivery} gauge; alert on staleness — a dead export and no traffic look identical in counter rates. + // Catches stalls with no failures to count (e.g. the flush trigger died): alert on staleness + // gated by offered/depth, so idle time isn't mistaken for a dead export. Gauge.builder( "billing.s3.last_delivery.epoch_seconds", lastDeliveryEpochSeconds, From df72274c8962c35242d071ce86b9dd40c1a25186 Mon Sep 17 00:00:00 2001 From: Hazel Date: Wed, 15 Jul 2026 10:42:43 -0700 Subject: [PATCH 29/65] Add BillingQueueTest --- .../service/provider/BillingQueueTest.java | 90 +++++++++++++++++++ 1 file changed, 90 insertions(+) create mode 100644 src/test/java/io/stargate/sgv2/jsonapi/service/provider/BillingQueueTest.java diff --git a/src/test/java/io/stargate/sgv2/jsonapi/service/provider/BillingQueueTest.java b/src/test/java/io/stargate/sgv2/jsonapi/service/provider/BillingQueueTest.java new file mode 100644 index 0000000000..91d98bfb2d --- /dev/null +++ b/src/test/java/io/stargate/sgv2/jsonapi/service/provider/BillingQueueTest.java @@ -0,0 +1,90 @@ +package io.stargate.sgv2.jsonapi.service.provider; + +import static org.assertj.core.api.Assertions.assertThat; + +import java.time.Instant; +import org.junit.jupiter.api.Test; + +/** Unit tests for {@link BillingQueue}: seal thresholds, drain limits, and batch metadata. */ +class BillingQueueTest { + + private static final Instant T0 = Instant.parse("2026-05-20T14:23:11Z"); + + @Test + void sealsByEventCount() { + var queue = new BillingQueue(2, 1_000_000, 10); + + queue.offer(T0, "a"); + assertThat(queue.shouldFlush()).isFalse(); + queue.offer(T0, "b"); + assertThat(queue.shouldFlush()).isTrue(); + } + + @Test + void sealsByBufferedBytes() { + // Each line counts as length + 1 (newline): "aaaa" = 5 bytes. + var queue = new BillingQueue(100, 10, 10); + + queue.offer(T0, "aaaa"); + assertThat(queue.shouldFlush()).isFalse(); + queue.offer(T0, "bbbb"); + assertThat(queue.shouldFlush()).isTrue(); + } + + @Test + void drainStopsAtMaxEventsAndLeavesTheRemainder() { + var queue = new BillingQueue(2, 1_000_000, 10); + queue.offer(T0, "a"); + queue.offer(T0, "b"); + queue.offer(T0, "c"); + + assertThat(queue.drain().lines()).containsExactly("a", "b"); + assertThat(queue.drain().lines()).containsExactly("c"); + assertThat(queue.drain().isEmpty()).isTrue(); + } + + @Test + void drainStopsAtMaxBytesAndLeavesTheRemainder() { + var queue = new BillingQueue(100, 10, 10); + queue.offer(T0, "aaaa"); + queue.offer(T0, "bbbb"); + queue.offer(T0, "cccc"); + + assertThat(queue.drain().lines()).containsExactly("aaaa", "bbbb"); + assertThat(queue.drain().lines()).containsExactly("cccc"); + assertThat(queue.drain().isEmpty()).isTrue(); + } + + @Test + void oldestEventAtIsTheMinimumAcrossTheBatchNotTheHead() { + var queue = new BillingQueue(10, 1_000_000, 10); + // Concurrent publishes can enqueue out of event-time order; the head is not the oldest. + queue.offer(T0.plusSeconds(5), "enqueued-first-but-newer"); + queue.offer(T0, "enqueued-second-but-older"); + + assertThat(queue.drain().oldestEventAt()).isEqualTo(T0); + } + + @Test + void offerRejectsWhenFull() { + var queue = new BillingQueue(10, 1_000_000, 2); + + assertThat(queue.offer(T0, "a")).isTrue(); + assertThat(queue.offer(T0, "b")).isTrue(); + assertThat(queue.offer(T0, "c")).isFalse(); + assertThat(queue.size()).isEqualTo(2); + } + + @Test + void byteSealResetsOnceDrained() { + var queue = new BillingQueue(100, 10, 10); + queue.offer(T0, "aaaa"); + queue.offer(T0, "bbbb"); + assertThat(queue.shouldFlush()).isTrue(); + + queue.drain(); + + assertThat(queue.isEmpty()).isTrue(); + assertThat(queue.shouldFlush()).isFalse(); // queuedBytes went back down with the drain + } +} From 5c8010e9a2f3eb5c69d1802634e116c5df6fe474 Mon Sep 17 00:00:00 2001 From: Hazel Date: Wed, 15 Jul 2026 10:54:12 -0700 Subject: [PATCH 30/65] Add S3BatchUploaderTest --- .../service/provider/S3BatchUploaderTest.java | 139 ++++++++++++++++++ 1 file changed, 139 insertions(+) create mode 100644 src/test/java/io/stargate/sgv2/jsonapi/service/provider/S3BatchUploaderTest.java diff --git a/src/test/java/io/stargate/sgv2/jsonapi/service/provider/S3BatchUploaderTest.java b/src/test/java/io/stargate/sgv2/jsonapi/service/provider/S3BatchUploaderTest.java new file mode 100644 index 0000000000..aa010bea7f --- /dev/null +++ b/src/test/java/io/stargate/sgv2/jsonapi/service/provider/S3BatchUploaderTest.java @@ -0,0 +1,139 @@ +package io.stargate.sgv2.jsonapi.service.provider; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import java.nio.charset.StandardCharsets; +import java.time.Duration; +import java.time.Instant; +import java.util.List; +import java.util.Optional; +import java.util.UUID; +import java.util.concurrent.CompletableFuture; +import java.util.regex.Pattern; +import org.junit.jupiter.api.Test; +import org.mockito.ArgumentCaptor; +import software.amazon.awssdk.core.async.AsyncRequestBody; +import software.amazon.awssdk.services.s3.S3AsyncClient; +import software.amazon.awssdk.services.s3.model.PutObjectRequest; +import software.amazon.awssdk.services.s3.model.PutObjectResponse; + +/** + * The S3 client is mocked; retries and per-call timeouts live in the client configuration, so + * exactly one {@code putObject} per upload is expected here. Real I/O is covered by {@code + * BillingS3ExportIntegrationTest}. + */ +class S3BatchUploaderTest { + + private static final Duration AWAIT = Duration.ofSeconds(5); + private static final Pattern KEY_PATTERN = + Pattern.compile("billing-events/2026/05/20/14/23/[0-9a-f-]{36}\\.jsonl"); + private static final String LINE_A = "{\"a\":1}"; + private static final String LINE_B = "{\"b\":2}"; + private static final BillingQueue.Batch BATCH = + new BillingQueue.Batch(List.of(LINE_A, LINE_B), Instant.parse("2026-05-20T14:23:11.482Z")); + + private static S3BatchUploader uploader(S3AsyncClient client) { + return new S3BatchUploader(client, "my-bucket"); + } + + private static CompletableFuture ok() { + return CompletableFuture.completedFuture(PutObjectResponse.builder().build()); + } + + // ============================================================ + // object layout — key and body + // ============================================================ + + @Test + void objectKeyUsesPathPrefixAndUtcMinutePathFromTimestamp() { + var id = UUID.fromString("8c0e9b8a-1d3a-4f6b-9c0d-1234567890ab"); + var key = S3BatchUploader.objectKey(Instant.parse("2026-05-20T14:23:11.482Z"), id); + assertThat(key) + .isEqualTo("billing-events/2026/05/20/14/23/8c0e9b8a-1d3a-4f6b-9c0d-1234567890ab.jsonl"); + } + + @Test + void toNdjsonJoinsLinesVerbatimWithTrailingNewlines() { + assertThat(S3BatchUploader.toNdjson(List.of(LINE_A, LINE_B))) + .isEqualTo((LINE_A + "\n" + LINE_B + "\n").getBytes(StandardCharsets.UTF_8)); + } + + @Test + void putsTheNdjsonBodyAtATimePartitionedKey() { + S3AsyncClient client = mock(S3AsyncClient.class); + when(client.putObject(any(PutObjectRequest.class), any(AsyncRequestBody.class))) + .thenReturn(ok()); + + uploader(client).upload(BATCH).await().atMost(AWAIT); + + var req = ArgumentCaptor.forClass(PutObjectRequest.class); + var body = ArgumentCaptor.forClass(AsyncRequestBody.class); + verify(client).putObject(req.capture(), body.capture()); + assertThat(req.getValue().bucket()).isEqualTo("my-bucket"); + // The key's minute path comes from the batch's oldestEventAt, not the wall clock. + assertThat(req.getValue().key()).matches(KEY_PATTERN.pattern()); + assertThat(req.getValue().contentType()).isEqualTo("application/x-ndjson"); + assertThat(body.getValue().contentLength()) + .hasValue((long) (LINE_A + "\n" + LINE_B + "\n").getBytes(StandardCharsets.UTF_8).length); + } + + @Test + void eachUploadGetsAFreshKey() { + S3AsyncClient client = mock(S3AsyncClient.class); + when(client.putObject(any(PutObjectRequest.class), any(AsyncRequestBody.class))) + .thenReturn(ok()); + + var uploader = uploader(client); + uploader.upload(BATCH).await().atMost(AWAIT); + uploader.upload(BATCH).await().atMost(AWAIT); + + var req = ArgumentCaptor.forClass(PutObjectRequest.class); + verify(client, times(2)).putObject(req.capture(), any(AsyncRequestBody.class)); + assertThat(req.getAllValues().stream().map(PutObjectRequest::key)).doesNotHaveDuplicates(); + } + + // ============================================================ + // failure — surfaces once; retries belong to the SDK client + // ============================================================ + + @Test + void uploadFailurePropagatesAndPutsExactlyOnceAtThisLayer() { + S3AsyncClient client = mock(S3AsyncClient.class); + when(client.putObject(any(PutObjectRequest.class), any(AsyncRequestBody.class))) + .thenReturn(CompletableFuture.failedFuture(new RuntimeException("simulated S3 failure"))); + + var uploader = uploader(client); + assertThatThrownBy(() -> uploader.upload(BATCH).await().atMost(AWAIT)) + .hasMessageContaining("simulated S3 failure"); + + // retries and per-call timeouts live in the client configuration, so exactly one putObject per upload is expected here + verify(client, times(1)).putObject(any(PutObjectRequest.class), any(AsyncRequestBody.class)); + } + + // ============================================================ + // lifecycle + config validation + // ============================================================ + + @Test + void closeClosesTheClient() { + S3AsyncClient client = mock(S3AsyncClient.class); + uploader(client).close(); + verify(client).close(); + } + + @Test + void createRejectsMissingRegionOrBucket() { + assertThatThrownBy(() -> S3BatchUploader.create(" ", "bucket", Optional.empty())) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("bucket-region"); + assertThatThrownBy(() -> S3BatchUploader.create("us-east-1", null, Optional.empty())) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("billing.s3.bucket"); + } +} From 8c26fc4eda320ba8b39d474a15bd692c27f2afce Mon Sep 17 00:00:00 2001 From: Hazel Date: Wed, 15 Jul 2026 11:07:11 -0700 Subject: [PATCH 31/65] Add BillingS3HandlerInstallerTest --- .../BillingS3HandlerInstallerTest.java | 83 +++++++++++++++++++ 1 file changed, 83 insertions(+) create mode 100644 src/test/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3HandlerInstallerTest.java diff --git a/src/test/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3HandlerInstallerTest.java b/src/test/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3HandlerInstallerTest.java new file mode 100644 index 0000000000..f085830f6f --- /dev/null +++ b/src/test/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3HandlerInstallerTest.java @@ -0,0 +1,83 @@ +package io.stargate.sgv2.jsonapi.service.provider; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import io.micrometer.core.instrument.simple.SimpleMeterRegistry; +import io.quarkus.runtime.ShutdownEvent; +import io.quarkus.runtime.StartupEvent; +import io.stargate.sgv2.jsonapi.config.BillingS3ExportConfig; +import java.time.Duration; +import java.util.Arrays; +import java.util.Optional; +import java.util.logging.Logger; +import org.junit.jupiter.api.Test; + +/** + * Unit tests for {@link BillingS3HandlerInstaller}: install/uninstall symmetry on the {@code + * billing.events} JUL logger, the disabled path, and fail-loud startup on bad config. Delivery + * through an installed handler is covered by {@code BillingS3ExportIntegrationTest}. + */ +class BillingS3HandlerInstallerTest { + + private static BillingS3ExportConfig config(boolean enabled, String bucket, String region) { + BillingS3ExportConfig config = mock(BillingS3ExportConfig.class); + when(config.enabled()).thenReturn(enabled); + when(config.bucket()).thenReturn(Optional.ofNullable(bucket)); + when(config.bucketRegion()).thenReturn(Optional.ofNullable(region)); + when(config.endpointOverride()).thenReturn(Optional.empty()); + when(config.maxEvents()).thenReturn(50); + when(config.maxBytes()).thenReturn(2_097_152L); + when(config.maxAge()).thenReturn(Duration.ofSeconds(30)); + when(config.queueCapacity()).thenReturn(100); + when(config.uploadConcurrency()).thenReturn(2); + when(config.shutdownTimeout()).thenReturn(Duration.ofSeconds(1)); + return config; + } + + private static long installedHandlers() { + return Arrays.stream( + Logger.getLogger(BillingS3HandlerInstaller.BILLING_LOGGER_NAME).getHandlers()) + .filter(BillingS3LogHandler.class::isInstance) + .count(); + } + + @Test + void disabledConfigInstallsNothing() { + var installer = + new BillingS3HandlerInstaller(config(false, null, null), new SimpleMeterRegistry()); + + installer.onStart(new StartupEvent()); + + assertThat(installedHandlers()).isZero(); + installer.onStop(new ShutdownEvent()); // must be a safe no-op without an installed handler + } + + @Test + void missingBucketFailsStartupLoudly() { + var installer = + new BillingS3HandlerInstaller(config(true, null, "us-east-1"), new SimpleMeterRegistry()); + + assertThatThrownBy(() -> installer.onStart(new StartupEvent())) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("bucket"); + assertThat(installedHandlers()).isZero(); + } + + @Test + void installsOnStartupAndRemovesAndClosesOnShutdown() { + var installer = + new BillingS3HandlerInstaller( + config(true, "my-bucket", "us-east-1"), new SimpleMeterRegistry()); + + installer.onStart(new StartupEvent()); + try { + assertThat(installedHandlers()).isEqualTo(1); + } finally { + installer.onStop(new ShutdownEvent()); + } + assertThat(installedHandlers()).isZero(); + } +} From 718ded30ca434c0bcd18be965550071a32a70559 Mon Sep 17 00:00:00 2001 From: Hazel Date: Wed, 15 Jul 2026 11:11:28 -0700 Subject: [PATCH 32/65] Add BillingMetricsTest --- .../service/provider/BillingMetricsTest.java | 58 +++++++++++++++++++ 1 file changed, 58 insertions(+) create mode 100644 src/test/java/io/stargate/sgv2/jsonapi/service/provider/BillingMetricsTest.java diff --git a/src/test/java/io/stargate/sgv2/jsonapi/service/provider/BillingMetricsTest.java b/src/test/java/io/stargate/sgv2/jsonapi/service/provider/BillingMetricsTest.java new file mode 100644 index 0000000000..ade91aac03 --- /dev/null +++ b/src/test/java/io/stargate/sgv2/jsonapi/service/provider/BillingMetricsTest.java @@ -0,0 +1,58 @@ +package io.stargate.sgv2.jsonapi.service.provider; + +import static org.assertj.core.api.Assertions.assertThat; + +import io.micrometer.core.instrument.simple.SimpleMeterRegistry; +import java.time.Instant; +import java.util.concurrent.atomic.AtomicInteger; +import org.junit.jupiter.api.Test; + +/** Guards the meter names and tags — dashboards and alerts key on these exact series. */ +class BillingMetricsTest { + + @Test + void countersFlowToTheExpectedSeries() { + var registry = new SimpleMeterRegistry(); + var metrics = new BillingMetrics(registry, () -> 0, 100); + + metrics.recordOffered(); + metrics.recordDropped(); + metrics.recordAbandonedAtShutdown(3); + metrics.recordBatchDelivered(2); + metrics.recordBatchFailed(5); + + assertThat(registry.counter("billing.s3.events.offered").count()).isEqualTo(1.0); + assertThat(registry.counter("billing.s3.events.dropped", "reason", "capacity").count()) + .isEqualTo(1.0); + assertThat(registry.counter("billing.s3.events.dropped", "reason", "shutdown").count()) + .isEqualTo(3.0); + assertThat(registry.counter("billing.s3.events.flushed").count()).isEqualTo(2.0); + assertThat(registry.counter("billing.s3.batches.uploaded").count()).isEqualTo(1.0); + assertThat(registry.counter("billing.s3.events.failed").count()).isEqualTo(5.0); + assertThat(registry.counter("billing.s3.batches.failed").count()).isEqualTo(1.0); + } + + @Test + void depthGaugeReadsTheLiveSupplier() { + var registry = new SimpleMeterRegistry(); + var depth = new AtomicInteger(7); + new BillingMetrics(registry, depth::get, 100); + + assertThat(registry.get("billing.s3.queue.depth").gauge().value()).isEqualTo(7.0); + depth.set(11); + assertThat(registry.get("billing.s3.queue.depth").gauge().value()).isEqualTo(11.0); + } + + @Test + void deliveryHeartbeatAdvancesOnDeliveredBatches() { + var registry = new SimpleMeterRegistry(); + var metrics = new BillingMetrics(registry, () -> 0, 100); + var heartbeat = registry.get("billing.s3.last_delivery.epoch_seconds").gauge(); + + assertThat(heartbeat.value()).isZero(); // never delivered + + long before = Instant.now().getEpochSecond(); + metrics.recordBatchDelivered(1); + assertThat(heartbeat.value()).isGreaterThanOrEqualTo(before); + } +} From 113de68fa2c90bea72dfb05c6ad81322a16db55f Mon Sep 17 00:00:00 2001 From: Hazel Date: Wed, 15 Jul 2026 11:14:12 -0700 Subject: [PATCH 33/65] format --- .../sgv2/jsonapi/service/provider/S3BatchUploaderTest.java | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/test/java/io/stargate/sgv2/jsonapi/service/provider/S3BatchUploaderTest.java b/src/test/java/io/stargate/sgv2/jsonapi/service/provider/S3BatchUploaderTest.java index aa010bea7f..90e8405354 100644 --- a/src/test/java/io/stargate/sgv2/jsonapi/service/provider/S3BatchUploaderTest.java +++ b/src/test/java/io/stargate/sgv2/jsonapi/service/provider/S3BatchUploaderTest.java @@ -112,7 +112,8 @@ void uploadFailurePropagatesAndPutsExactlyOnceAtThisLayer() { assertThatThrownBy(() -> uploader.upload(BATCH).await().atMost(AWAIT)) .hasMessageContaining("simulated S3 failure"); - // retries and per-call timeouts live in the client configuration, so exactly one putObject per upload is expected here + // retries and per-call timeouts live in the client configuration, so exactly one putObject per + // upload is expected here verify(client, times(1)).putObject(any(PutObjectRequest.class), any(AsyncRequestBody.class)); } From 2818ff4a0c1f3b1d166867fef54d6e933a9259aa Mon Sep 17 00:00:00 2001 From: Hazel Date: Wed, 15 Jul 2026 14:10:57 -0700 Subject: [PATCH 34/65] update BillingQueue java doc --- .../stargate/sgv2/jsonapi/service/provider/BillingQueue.java | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingQueue.java b/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingQueue.java index 873e0df20d..b016b170f9 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingQueue.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingQueue.java @@ -9,8 +9,9 @@ /** * Bounded in-memory buffer that owns the batching policy of the billing S3 export: it decides when - * a batch is sealed ({@code maxEvents} line-count seal, {@code maxBytes} body-size seal) and hands - * out drained {@link Batch}es. + * a batch is sealed ({@code maxEvents} lines or {@code maxBytes} UTF-8 NDJSON bytes) and hands out + * drained {@link Batch}es. A batch may exceed {@code maxBytes} by one whole line; lines are never + * split. */ public final class BillingQueue { From 1977782ad1d3030790bb056639f9c0f660396b6e Mon Sep 17 00:00:00 2001 From: Hazel Date: Wed, 15 Jul 2026 14:27:31 -0700 Subject: [PATCH 35/65] update BillingS3ExportConfig java doc --- .../io/stargate/sgv2/jsonapi/config/BillingS3ExportConfig.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/main/java/io/stargate/sgv2/jsonapi/config/BillingS3ExportConfig.java b/src/main/java/io/stargate/sgv2/jsonapi/config/BillingS3ExportConfig.java index f8b927477f..28dd29e023 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/config/BillingS3ExportConfig.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/config/BillingS3ExportConfig.java @@ -26,7 +26,7 @@ public interface BillingS3ExportConfig { @WithDefault("50") int maxEvents(); - /** Byte-size seal on the buffered NDJSON body. */ + /** UTF-8 NDJSON byte seal; a batch may exceed it by one whole event. */ @WithDefault("2097152") long maxBytes(); From 4c4f5036f878ffb7a389aaddceb0ac0990061afd Mon Sep 17 00:00:00 2001 From: Hazel Date: Wed, 15 Jul 2026 14:31:59 -0700 Subject: [PATCH 36/65] update BillingMetrics java doc --- .../jsonapi/service/provider/BillingMetrics.java | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingMetrics.java b/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingMetrics.java index 88d30636ae..e2dc493811 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingMetrics.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingMetrics.java @@ -11,10 +11,14 @@ import org.slf4j.LoggerFactory; /** - * Meters for the billing S3 export. Every offered event ends in exactly one terminal counter — - * {@code flushed}, {@code failed}, or {@code dropped(reason=capacity|shutdown)} — or is still - * visible in {@code queue.depth} / an in-flight batch, so loss is always accounted for. {@code - * last_delivery.epoch_seconds} is the freshness heartbeat to alert on. + * During normal operation, {@code offered = flushed + failed + dropped(capacity) + queue.depth + + * events in in-flight batches}. + * + *

At shutdown, buffered events left after the drain budget are added to {@code + * dropped(shutdown)}. Final counters can be lower than {@code offered} if a concurrent publish + * misses the final queue snapshot or an in-flight upload does not settle before process exit. + * + *

{@code last_delivery.epoch_seconds} is the delivery heartbeat. */ public final class BillingMetrics { @@ -85,14 +89,14 @@ public void recordAbandonedAtShutdown(int size) { droppedShutdown.increment(size); } - /** A batch of events lines landed in S3; bumps the delivery heartbeat. */ + /** A batch of event lines landed in S3; bumps the delivery heartbeat. */ public void recordBatchDelivered(int size) { flushed.increment(size); batchesUploaded.increment(); lastDeliveryEpochSeconds.set(Instant.now().getEpochSecond()); } - /** A batch of events lines was given up after the uploader exhausted its retries. */ + /** A batch of event lines was given up after the uploader exhausted its retries. */ public void recordBatchFailed(int size) { failed.increment(size); batchesFailed.increment(); From 4c9e54046e808841bfe85d50b1453d9f1b84d7ab Mon Sep 17 00:00:00 2001 From: Hazel Date: Wed, 15 Jul 2026 14:48:13 -0700 Subject: [PATCH 37/65] update BillingS3LogHandler --- .../service/provider/BillingS3LogHandler.java | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandler.java b/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandler.java index a7e40abd10..55f9df630e 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandler.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandler.java @@ -22,9 +22,8 @@ * decides what an S3 object looks like, and this class decides when uploads run — the flush * triggers (seal on publish, age tick, drain on close), the upload-concurrency gate, and metrics. * - *

Delivery is at-most-once by design: publish never blocks and never throws — when the buffer is - * full, lines are dropped and counted — and close() gives up loudly once {@code shutdownTimeout} is - * exhausted. + *

Delivery is at-most-once by design: publish never waits for queue capacity, full buffers drop + * new lines, and close drains best-effort within {@code shutdownTimeout}. */ public final class BillingS3LogHandler extends Handler { @@ -179,7 +178,7 @@ private void maybeFlush() { *

Catches everything: an escaped throwable would silently cancel all future runs of a * fixed-rate task. */ - private void onAgeTick() { + void onAgeTick() { try { if (!buffer.isEmpty()) { tryFlush(); @@ -218,12 +217,14 @@ private Uni uploadBatch(BillingQueue.Batch batch) { return Uni.createFrom().voidItem(); } int size = batch.size(); - return uploader - .upload(batch) + // Runs immediately on subscription; deferred only turns a throw before upload() returns a Uni + // into a Uni failure handled below. + return Uni.createFrom() + .deferred(() -> uploader.upload(batch)) .onItem() .invoke(() -> metrics.recordBatchDelivered(size)) .onFailure() - .invoke(t -> LOG.error("Failed to upload billing S3 batch ({} size)", size, t)) + .invoke(t -> LOG.error("Failed to upload billing S3 batch ({} events)", size, t)) .onFailure() .recoverWithItem( () -> { From 503072cef89697330f6848e2355fca564c299621 Mon Sep 17 00:00:00 2001 From: Hazel Date: Wed, 15 Jul 2026 17:02:54 -0700 Subject: [PATCH 38/65] Add @VisibleForTesting --- .../service/provider/BillingQueue.java | 7 ++ .../service/provider/BillingS3LogHandler.java | 1 + .../service/provider/BillingQueueTest.java | 72 +++++++++++++++++++ 3 files changed, 80 insertions(+) diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingQueue.java b/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingQueue.java index b016b170f9..33171df5c9 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingQueue.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingQueue.java @@ -1,5 +1,6 @@ package io.stargate.sgv2.jsonapi.service.provider; +import com.google.common.annotations.VisibleForTesting; import java.time.Instant; import java.util.ArrayList; import java.util.List; @@ -65,6 +66,12 @@ public int size() { return queue.size(); } + /** Buffered-bytes counter, exposed for accounting assertions only (no production caller). */ + @VisibleForTesting + long queuedBytes() { + return queuedBytes.get(); + } + /** Removes and returns up to one sealed batch (possibly partial, possibly {@code EMPTY}). */ public Batch drain() { List lines = new ArrayList<>(batchSize); diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandler.java b/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandler.java index 55f9df630e..97af685d37 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandler.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandler.java @@ -178,6 +178,7 @@ private void maybeFlush() { *

Catches everything: an escaped throwable would silently cancel all future runs of a * fixed-rate task. */ + @VisibleForTesting void onAgeTick() { try { if (!buffer.isEmpty()) { diff --git a/src/test/java/io/stargate/sgv2/jsonapi/service/provider/BillingQueueTest.java b/src/test/java/io/stargate/sgv2/jsonapi/service/provider/BillingQueueTest.java index 91d98bfb2d..5d0cd55f5a 100644 --- a/src/test/java/io/stargate/sgv2/jsonapi/service/provider/BillingQueueTest.java +++ b/src/test/java/io/stargate/sgv2/jsonapi/service/provider/BillingQueueTest.java @@ -3,6 +3,17 @@ import static org.assertj.core.api.Assertions.assertThat; import java.time.Instant; +import java.util.ArrayList; +import java.util.HashSet; +import java.util.List; +import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; import org.junit.jupiter.api.Test; /** Unit tests for {@link BillingQueue}: seal thresholds, drain limits, and batch metadata. */ @@ -75,6 +86,67 @@ void offerRejectsWhenFull() { assertThat(queue.size()).isEqualTo(2); } + @Test + void concurrentOfferAndDrainKeepsAccountingConsistent() throws Exception { + var queue = new BillingQueue(10, 1_000_000, 5_000); + int threads = 4; + int perThread = 500; + Set published = ConcurrentHashMap.newKeySet(); + List drained = new ArrayList<>(); // touched only by the drainer thread until join + + AtomicBoolean producersDone = new AtomicBoolean(false); + Thread drainer = + new Thread( + () -> { + while (!producersDone.get() || !queue.isEmpty()) { + var batch = queue.drain(); + if (batch.isEmpty()) { + Thread.onSpinWait(); + } else { + drained.addAll(batch.lines()); + } + } + }, + "billing-queue-test-drainer"); + drainer.start(); + + ExecutorService executor = Executors.newFixedThreadPool(threads); + try { + CountDownLatch start = new CountDownLatch(1); + List> futures = new ArrayList<>(); + for (int t = 0; t < threads; t++) { + int threadId = t; + futures.add( + executor.submit( + () -> { + start.await(); + for (int i = 0; i < perThread; i++) { + String line = "{\"t\":" + threadId + ",\"i\":" + i + "}"; + published.add(line); + assertThat(queue.offer(T0, line)).isTrue(); // capacity is never reached + } + return null; + })); + } + start.countDown(); + for (Future future : futures) { + future.get(30, TimeUnit.SECONDS); + } + } finally { + executor.shutdown(); + } + producersDone.set(true); + drainer.join(TimeUnit.SECONDS.toMillis(30)); + assertThat(drainer.isAlive()).isFalse(); + + // Every offered line is drained exactly once, and the byte accounting lands back on zero: + // concurrent add/subtract may transiently disagree, but the settled state must not drift. + assertThat(drained).hasSize(threads * perThread); + assertThat(new HashSet<>(drained)).isEqualTo(published); + assertThat(queue.isEmpty()).isTrue(); + assertThat(queue.queuedBytes()).isZero(); + } + @Test void byteSealResetsOnceDrained() { var queue = new BillingQueue(100, 10, 10); From 7ae577dc6830b8ff2b2b7e9a5cf90003c0365d16 Mon Sep 17 00:00:00 2001 From: Hazel Date: Wed, 15 Jul 2026 17:06:05 -0700 Subject: [PATCH 39/65] Add BillingS3LogHandler unit test --- .../provider/BillingS3LogHandlerTest.java | 707 ++++++++++++++++++ 1 file changed, 707 insertions(+) create mode 100644 src/test/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandlerTest.java diff --git a/src/test/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandlerTest.java b/src/test/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandlerTest.java new file mode 100644 index 0000000000..6e43aff18c --- /dev/null +++ b/src/test/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandlerTest.java @@ -0,0 +1,707 @@ +package io.stargate.sgv2.jsonapi.service.provider; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.awaitility.Awaitility.await; + +import io.micrometer.core.instrument.simple.SimpleMeterRegistry; +import io.smallrye.mutiny.Uni; +import io.smallrye.mutiny.infrastructure.Infrastructure; +import java.time.Duration; +import java.time.Instant; +import java.util.ArrayList; +import java.util.HashSet; +import java.util.List; +import java.util.Set; +import java.util.concurrent.BlockingQueue; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.LinkedBlockingQueue; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.logging.Level; +import java.util.logging.LogRecord; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; + +/** + * Unit tests for {@link BillingS3LogHandler}: the flush triggers (seal on publish, age tick, drain + * on close), the upload-concurrency gate, failure containment, and the at-most-once accounting + * invariants under concurrent publish. The uploader is a programmable in-memory fake; real S3 I/O + * is covered by {@code BillingS3ExportIntegrationTest}. + * + *

Ordering is only asserted in single-producer tests (the buffer is FIFO for a single thread); + * concurrency tests assert set-equality and counter reconciliation, never interleaving order. + */ +class BillingS3LogHandlerTest { + + private static final Duration AWAIT = Duration.ofSeconds(10); + private static final Duration NEVER = Duration.ofHours(1); + private static final Instant T0 = Instant.parse("2026-05-20T14:23:11Z"); + + /** + * First-ever Uni creation in a JVM registers the SmallRye context-propagation provider through + * {@code ContextManagerProvider.instance()}, whose ServiceLoader loop both CAS-races concurrent + * callers and throws "ContextManagerProvider already set" when it discovers a second provider — + * possibly after having registered the first. Racing that from concurrent producer threads makes + * publish() throw. Quarkus registers the provider single-threaded at boot, so only this + * bare-JUnit JVM needs the deterministic warm-up. + */ + @BeforeAll + static void warmUpMutinyInfrastructure() { + try { + io.smallrye.context.SmallRyeContextManagerProvider.getManager(); + } catch (IllegalStateException alreadySetOrDuplicate) { + // The provider is registered even when the duplicate-discovery branch throws; either way + // ContextManagerProvider.INSTANCE is now set and concurrent callers can no longer race it. + } + Uni.createFrom() + .item(0) + .runSubscriptionOn(Infrastructure.getDefaultWorkerPool()) + .await() + .atMost(AWAIT); + } + + private static final String OFFERED = "billing.s3.events.offered"; + private static final String FLUSHED = "billing.s3.events.flushed"; + private static final String EVENTS_FAILED = "billing.s3.events.failed"; + private static final String BATCHES_UPLOADED = "billing.s3.batches.uploaded"; + private static final String BATCHES_FAILED = "billing.s3.batches.failed"; + private static final String DROPPED = "billing.s3.events.dropped"; + + // ============================================================ + // Fake uploader + // ============================================================ + + /** + * Programmable {@link BillingS3LogHandler.AsyncBatchUploader}: records every batch and settles + * the returned Uni per {@link Mode}. Never blocks a caller thread — HOLD parks the completion in + * {@code held} for the test to release explicitly. + */ + static final class RecordingUploader implements BillingS3LogHandler.AsyncBatchUploader { + enum Mode { + COMPLETE, + HOLD, + FAIL, + THROW_SYNC + } + + volatile Mode mode = Mode.COMPLETE; + final List batches = new CopyOnWriteArrayList<>(); + final BlockingQueue> held = new LinkedBlockingQueue<>(); + final AtomicInteger inFlight = new AtomicInteger(); + final AtomicInteger maxInFlight = new AtomicInteger(); + volatile boolean closed; + + @Override + public Uni upload(BillingQueue.Batch batch) { + batches.add(batch); + if (mode == Mode.THROW_SYNC) { + throw new RuntimeException("simulated synchronous uploader failure"); + } + int now = inFlight.incrementAndGet(); + maxInFlight.accumulateAndGet(now, Math::max); + CompletableFuture future = new CompletableFuture<>(); + future.whenComplete((v, t) -> inFlight.decrementAndGet()); + switch (mode) { + case COMPLETE -> future.complete(null); + case FAIL -> future.completeExceptionally(new RuntimeException("simulated upload failure")); + case HOLD -> held.add(future); + default -> throw new IllegalStateException("unexpected mode " + mode); + } + return Uni.createFrom().completionStage(future); + } + + /** Completes one held upload, waiting for it to exist first. */ + void releaseOne() throws InterruptedException { + CompletableFuture future = held.poll(AWAIT.toSeconds(), TimeUnit.SECONDS); + assertThat(future).as("a held upload to release").isNotNull(); + future.complete(null); + } + + /** Switches to pass-through and completes everything currently held. */ + void releaseAllAndComplete() { + mode = Mode.COMPLETE; + CompletableFuture future; + while ((future = held.poll()) != null) { + future.complete(null); + } + } + + List allLines() { + return batches.stream().flatMap(b -> b.lines().stream()).toList(); + } + + @Override + public void close() { + closed = true; + } + } + + // ============================================================ + // Helpers + // ============================================================ + + private static BillingS3LogHandler newHandler( + RecordingUploader uploader, + SimpleMeterRegistry registry, + int maxEvents, + long maxBytes, + int queueCapacity, + int uploadConcurrency) { + return newHandler( + uploader, + registry, + maxEvents, + maxBytes, + queueCapacity, + uploadConcurrency, + Duration.ofSeconds(5)); + } + + private static BillingS3LogHandler newHandler( + RecordingUploader uploader, + SimpleMeterRegistry registry, + int maxEvents, + long maxBytes, + int queueCapacity, + int uploadConcurrency, + Duration shutdownTimeout) { + return new BillingS3LogHandler( + uploader, + registry, + maxEvents, + maxBytes, + NEVER, + queueCapacity, + uploadConcurrency, + shutdownTimeout); + } + + private static LogRecord record(String message) { + return record(T0, message); + } + + private static LogRecord record(Instant at, String message) { + LogRecord logRecord = new LogRecord(Level.INFO, message); + logRecord.setInstant(at); + return logRecord; + } + + private static double counter(SimpleMeterRegistry registry, String name, String... tags) { + return registry.counter(name, tags).count(); + } + + // ============================================================ + // Behavior — publish and flush triggers + // ============================================================ + + @Test + void publishIgnoresNullRecordAndBlankLines() { + var uploader = new RecordingUploader(); + var registry = new SimpleMeterRegistry(); + var handler = newHandler(uploader, registry, 1, 1_000_000, 10, 1); + try { + handler.publish(null); + handler.publish(record(null)); + handler.publish(record(" ")); + + assertThat(uploader.batches).isEmpty(); + assertThat(counter(registry, OFFERED)).isZero(); + } finally { + handler.close(); + } + } + + @Test + void sealsByCountAndShipsExactBatch() { + var uploader = new RecordingUploader(); + var registry = new SimpleMeterRegistry(); + var handler = newHandler(uploader, registry, 3, 1_000_000, 10, 2); + try { + // Enqueue order is publish order for a single producer; event-time order is not (the second + // record is older on purpose, so oldestEventAt must be the min, not the head). + handler.publish(record(T0.plusSeconds(5), "{\"e\":1}")); + handler.publish(record(T0, "{\"e\":2}")); + await() + .during(Duration.ofMillis(200)) + .atMost(Duration.ofSeconds(2)) + .until(() -> uploader.batches.isEmpty()); + + handler.publish(record(T0.plusSeconds(9), "{\"e\":3}")); + + await().atMost(AWAIT).untilAsserted(() -> assertThat(uploader.batches).hasSize(1)); + var batch = uploader.batches.get(0); + assertThat(batch.lines()).containsExactly("{\"e\":1}", "{\"e\":2}", "{\"e\":3}"); + assertThat(batch.oldestEventAt()).isEqualTo(T0); + await() + .atMost(AWAIT) + .untilAsserted( + () -> { + assertThat(counter(registry, OFFERED)).isEqualTo(3.0); + assertThat(counter(registry, FLUSHED)).isEqualTo(3.0); + assertThat(counter(registry, BATCHES_UPLOADED)).isEqualTo(1.0); + assertThat(counter(registry, DROPPED, "reason", "capacity")).isZero(); + }); + } finally { + handler.close(); + } + } + + @Test + void sealsByBufferedBytes() { + var uploader = new RecordingUploader(); + var registry = new SimpleMeterRegistry(); + // Lines count as length + 1: two 4-char lines hit the 10-byte seal together. + var handler = newHandler(uploader, registry, 100, 10, 10, 2); + try { + handler.publish(record("aaaa")); + handler.publish(record("bbbb")); + + await().atMost(AWAIT).untilAsserted(() -> assertThat(uploader.batches).hasSize(1)); + assertThat(uploader.batches.get(0).lines()).containsExactly("aaaa", "bbbb"); + } finally { + handler.close(); + } + } + + @Test + void noShipmentBelowSealUntilAgeTick() { + var uploader = new RecordingUploader(); + var registry = new SimpleMeterRegistry(); + var handler = newHandler(uploader, registry, 100, 1_000_000, 10, 2); + try { + handler.publish(record("{\"e\":1}")); + handler.publish(record("{\"e\":2}")); + await() + .during(Duration.ofMillis(200)) + .atMost(Duration.ofSeconds(2)) + .until(() -> uploader.batches.isEmpty()); + + // Deterministic age trigger: call the tick directly instead of waiting for the scheduler. + handler.onAgeTick(); + + await().atMost(AWAIT).untilAsserted(() -> assertThat(uploader.batches).hasSize(1)); + assertThat(uploader.batches.get(0).lines()).containsExactly("{\"e\":1}", "{\"e\":2}"); + } finally { + handler.close(); + } + } + + @Test + void ageTickIsScheduledForReal() { + var uploader = new RecordingUploader(); + var registry = new SimpleMeterRegistry(); + var handler = + new BillingS3LogHandler( + uploader, + registry, + 100, + 1_000_000, + Duration.ofMillis(100), + 10, + 2, + Duration.ofSeconds(5)); + try { + handler.publish(record("{\"e\":1}")); + + // No seal is reached; only the scheduled fixed-rate tick can ship this line. + await().atMost(AWAIT).untilAsserted(() -> assertThat(uploader.batches).isNotEmpty()); + assertThat(uploader.allLines()).containsExactly("{\"e\":1}"); + } finally { + handler.close(); + } + } + + // ============================================================ + // Behavior — failure containment + // ============================================================ + + @Test + void uploadFailureCountsBatchAndPipelineSurvives() { + var uploader = new RecordingUploader(); + var registry = new SimpleMeterRegistry(); + var handler = newHandler(uploader, registry, 2, 1_000_000, 10, 1); + try { + uploader.mode = RecordingUploader.Mode.FAIL; + handler.publish(record("{\"e\":1}")); + handler.publish(record("{\"e\":2}")); + + await() + .atMost(AWAIT) + .untilAsserted( + () -> { + assertThat(counter(registry, BATCHES_FAILED)).isEqualTo(1.0); + assertThat(counter(registry, EVENTS_FAILED)).isEqualTo(2.0); + }); + + // The failure released the in-flight slot: the next sealed batch still ships. + uploader.mode = RecordingUploader.Mode.COMPLETE; + handler.publish(record("{\"e\":3}")); + handler.publish(record("{\"e\":4}")); + + await() + .atMost(AWAIT) + .untilAsserted( + () -> { + assertThat(counter(registry, FLUSHED)).isEqualTo(2.0); + assertThat(counter(registry, BATCHES_UPLOADED)).isEqualTo(1.0); + }); + } finally { + handler.close(); + } + } + + @Test + void synchronousUploaderThrowIsContained() { + var uploader = new RecordingUploader(); + var registry = new SimpleMeterRegistry(); + var handler = newHandler(uploader, registry, 2, 1_000_000, 10, 1); + try { + uploader.mode = RecordingUploader.Mode.THROW_SYNC; + handler.publish(record("{\"e\":1}")); + handler.publish(record("{\"e\":2}")); + + await() + .atMost(AWAIT) + .untilAsserted( + () -> { + assertThat(counter(registry, BATCHES_FAILED)).isEqualTo(1.0); + assertThat(counter(registry, EVENTS_FAILED)).isEqualTo(2.0); + }); + + uploader.mode = RecordingUploader.Mode.COMPLETE; + handler.publish(record("{\"e\":3}")); + handler.publish(record("{\"e\":4}")); + + await() + .atMost(AWAIT) + .untilAsserted(() -> assertThat(counter(registry, FLUSHED)).isEqualTo(2.0)); + } finally { + handler.close(); + } + } + + // ============================================================ + // Behavior — close + // ============================================================ + + @Test + void closeDrainsRemainderAndClosesUploader() { + var uploader = new RecordingUploader(); + var registry = new SimpleMeterRegistry(); + var handler = newHandler(uploader, registry, 100, 1_000_000, 10, 2); + + handler.publish(record("{\"e\":1}")); + handler.publish(record("{\"e\":2}")); + handler.publish(record("{\"e\":3}")); + handler.close(); + + // close() is synchronous: by the time it returns the drain has settled and counted. + assertThat(uploader.allLines()).containsExactly("{\"e\":1}", "{\"e\":2}", "{\"e\":3}"); + assertThat(uploader.closed).isTrue(); + assertThat(counter(registry, FLUSHED)).isEqualTo(3.0); + assertThat(counter(registry, DROPPED, "reason", "shutdown")).isZero(); + } + + @Test + void closeTimeoutCountsAbandoned() throws Exception { + var uploader = new RecordingUploader(); + var registry = new SimpleMeterRegistry(); + var handler = newHandler(uploader, registry, 1, 1_000_000, 10, 1, Duration.ofMillis(200)); + try { + uploader.mode = RecordingUploader.Mode.HOLD; + handler.publish(record("{\"e\":1}")); + // Wait until the first batch is in flight (and held) so the queued remainder is exact. + await().atMost(AWAIT).untilAsserted(() -> assertThat(uploader.batches).hasSize(1)); + handler.publish(record("{\"e\":2}")); + handler.publish(record("{\"e\":3}")); + + long startNanos = System.nanoTime(); + handler.close(); + Duration elapsed = Duration.ofNanos(System.nanoTime() - startNanos); + + assertThat(elapsed).isLessThan(Duration.ofSeconds(3)); + assertThat(counter(registry, DROPPED, "reason", "shutdown")).isEqualTo(2.0); + assertThat(uploader.closed).isTrue(); + } finally { + uploader.releaseAllAndComplete(); + } + } + + @Test + void closeIsIdempotentAndPublishAfterCloseIsSafe() { + var uploader = new RecordingUploader(); + var registry = new SimpleMeterRegistry(); + var handler = newHandler(uploader, registry, 1, 1_000_000, 10, 1); + + handler.close(); + handler.close(); // JUL Handler.close() contract: idempotent + + // A racing thread may publish after close; it must never throw (JUL handler contract). + handler.publish(record("{\"late\":1}")); + assertThat(counter(registry, OFFERED)).isEqualTo(1.0); + } + + // ============================================================ + // Concurrency — invariant style, no interleaving assertions + // ============================================================ + + @Test + void multiProducerNoLossNoDuplication() throws Exception { + var uploader = new RecordingUploader(); + var registry = new SimpleMeterRegistry(); + var handler = newHandler(uploader, registry, 50, 1_000_000_000L, 10_000, 4); + + int threads = 8; + int perThread = 500; + Set published = ConcurrentHashMap.newKeySet(); + runProducers( + threads, + (threadId) -> { + for (int i = 0; i < perThread; i++) { + String line = "{\"t\":" + threadId + ",\"i\":" + i + "}"; + published.add(line); + handler.publish(record(line)); + } + }); + handler.close(); + + List delivered = uploader.allLines(); + assertThat(delivered).hasSize(threads * perThread); + assertThat(new HashSet<>(delivered)).isEqualTo(published); + assertThat(counter(registry, OFFERED)).isEqualTo(threads * perThread); + assertThat(counter(registry, FLUSHED)).isEqualTo(threads * perThread); + assertThat(counter(registry, DROPPED, "reason", "capacity")).isZero(); + assertThat(counter(registry, DROPPED, "reason", "shutdown")).isZero(); + } + + @Test + void overflowAccountingReconciles() throws Exception { + var uploader = new RecordingUploader(); + var registry = new SimpleMeterRegistry(); + // No seal is ever reached (count seal above capacity, byte seal huge): nothing drains while + // producers run, so every line beyond the 64-slot buffer is a deterministic capacity drop. + var handler = newHandler(uploader, registry, 1000, 1_000_000_000L, 64, 4); + + int threads = 4; + int perThread = 500; + Set published = ConcurrentHashMap.newKeySet(); + runProducers( + threads, + (threadId) -> { + for (int i = 0; i < perThread; i++) { + String line = "{\"t\":" + threadId + ",\"i\":" + i + "}"; + published.add(line); + handler.publish(record(line)); + } + }); + + uploader.releaseAllAndComplete(); + handler.close(); + + double offered = counter(registry, OFFERED); + double flushed = counter(registry, FLUSHED); + double droppedCapacity = counter(registry, DROPPED, "reason", "capacity"); + double droppedShutdown = counter(registry, DROPPED, "reason", "shutdown"); + assertThat(offered).isEqualTo(threads * perThread); + assertThat(flushed).isEqualTo(64.0); + assertThat(droppedCapacity).isEqualTo(threads * perThread - 64.0); + assertThat(flushed + droppedCapacity + droppedShutdown).isEqualTo(offered); + + List delivered = uploader.allLines(); + assertThat(delivered).doesNotHaveDuplicates(); + assertThat(published).containsAll(delivered); + } + + @Test + void publishNeverBlocksWhenUploaderStalls() throws Exception { + var uploader = new RecordingUploader(); + var registry = new SimpleMeterRegistry(); + var handler = newHandler(uploader, registry, 1, 1_000_000, 8, 1); + try { + uploader.mode = RecordingUploader.Mode.HOLD; + handler.publish(record("{\"i\":0}")); + // Wait for the single slot to be claimed and its 1-line batch drained: from here the queue + // is empty, the slot is stuck, and every subsequent count is deterministic. + await().atMost(AWAIT).untilAsserted(() -> assertThat(uploader.batches).hasSize(1)); + + long startNanos = System.nanoTime(); + for (int i = 1; i < 50; i++) { + handler.publish(record("{\"i\":" + i + "}")); + } + Duration elapsed = Duration.ofNanos(System.nanoTime() - startNanos); + + assertThat(elapsed).isLessThan(Duration.ofSeconds(2)); + assertThat(counter(registry, OFFERED)).isEqualTo(50.0); + // 1 in flight + 8 buffered; the other 41 dropped without ever blocking the caller. + assertThat(counter(registry, DROPPED, "reason", "capacity")).isEqualTo(41.0); + } finally { + uploader.releaseAllAndComplete(); + handler.close(); + } + } + + @Test + void concurrencyGateCapsParallelUploads() throws Exception { + var uploader = new RecordingUploader(); + var registry = new SimpleMeterRegistry(); + var handler = newHandler(uploader, registry, 1, 1_000_000, 100, 2); + try { + uploader.mode = RecordingUploader.Mode.HOLD; + for (int i = 0; i < 10; i++) { + handler.publish(record("{\"i\":" + i + "}")); + } + + // Both slots claim work; the rest stays queued behind the gate. + await().atMost(AWAIT).untilAsserted(() -> assertThat(uploader.inFlight.get()).isEqualTo(2)); + + // Each release lets exactly the next batch through, one at a time. + while (uploader.batches.size() < 10) { + uploader.releaseOne(); + int expected = Math.min(uploader.batches.size() + 1, 10); + await() + .atMost(AWAIT) + .untilAsserted( + () -> assertThat(uploader.batches.size()).isGreaterThanOrEqualTo(expected)); + } + + assertThat(uploader.maxInFlight.get()).isEqualTo(2); + assertThat(uploader.allLines()).hasSize(10).doesNotHaveDuplicates(); + } finally { + uploader.releaseAllAndComplete(); + handler.close(); + } + } + + @Test + void settledUploadChainsNextBatchWithoutNewPublish() throws Exception { + var uploader = new RecordingUploader(); + var registry = new SimpleMeterRegistry(); + var handler = newHandler(uploader, registry, 2, 1_000_000, 100, 1); + try { + uploader.mode = RecordingUploader.Mode.HOLD; + for (int i = 1; i <= 6; i++) { + handler.publish(record("{\"i\":" + i + "}")); + } + + // Single slot: exactly one upload starts, the two other sealed batches wait behind it. + await().atMost(AWAIT).untilAsserted(() -> assertThat(uploader.batches).hasSize(1)); + await() + .during(Duration.ofMillis(200)) + .atMost(Duration.ofSeconds(2)) + .until(() -> uploader.batches.size() == 1); + assertThat(uploader.batches.get(0).lines()).containsExactly("{\"i\":1}", "{\"i\":2}"); + + // No further publish happens: each settle must chain the next flush on its own. + uploader.releaseOne(); + await().atMost(AWAIT).untilAsserted(() -> assertThat(uploader.batches).hasSize(2)); + assertThat(uploader.batches.get(1).lines()).containsExactly("{\"i\":3}", "{\"i\":4}"); + + uploader.releaseOne(); + await().atMost(AWAIT).untilAsserted(() -> assertThat(uploader.batches).hasSize(3)); + assertThat(uploader.batches.get(2).lines()).containsExactly("{\"i\":5}", "{\"i\":6}"); + + uploader.releaseOne(); + } finally { + uploader.releaseAllAndComplete(); + handler.close(); + } + } + + @Test + void closeRacingProducersNeverHangsAndReconciles() throws Exception { + var uploader = new RecordingUploader(); + var registry = new SimpleMeterRegistry(); + var handler = newHandler(uploader, registry, 5, 1_000_000_000L, 1000, 4, Duration.ofSeconds(1)); + + int threads = 4; + // Producers run flat out until stopped; the cap only bounds memory and assertion cost, sized + // so publishing is still in full flight when close() lands ~50ms in. + int perThreadCap = 200_000; + Set published = ConcurrentHashMap.newKeySet(); + List producerErrors = new CopyOnWriteArrayList<>(); + AtomicBoolean stop = new AtomicBoolean(false); + ExecutorService executor = Executors.newFixedThreadPool(threads); + List> futures = new ArrayList<>(); + for (int t = 0; t < threads; t++) { + int threadId = t; + futures.add( + executor.submit( + () -> { + for (int i = 0; i < perThreadCap && !stop.get(); i++) { + String line = "{\"t\":" + threadId + ",\"i\":" + i + "}"; + published.add(line); + try { + handler.publish(record(line)); + } catch (Throwable error) { + producerErrors.add(error); + return; + } + } + })); + } + + Thread.sleep(50); // let producers overlap the close below + handler.close(); + stop.set(true); + for (Future future : futures) { + future.get(AWAIT.toSeconds(), TimeUnit.SECONDS); + } + executor.shutdown(); + + // publish must never throw, close must return, and the books must stay consistent — lines + // published after close may be dropped without being counted, so the reconciliation is <=. + // Plain java.util.Set operations keep these checks O(n); AssertJ's containsAll would scan. + assertThat(producerErrors).isEmpty(); + List delivered = uploader.allLines(); + Set deliveredSet = new HashSet<>(delivered); + assertThat(deliveredSet).as("delivered lines must not repeat").hasSize(delivered.size()); + assertThat(published.containsAll(deliveredSet)) + .as("every delivered line must have been published") + .isTrue(); + double flushed = counter(registry, FLUSHED); + double droppedCapacity = counter(registry, DROPPED, "reason", "capacity"); + double droppedShutdown = counter(registry, DROPPED, "reason", "shutdown"); + assertThat(flushed + droppedCapacity + droppedShutdown) + .isLessThanOrEqualTo(counter(registry, OFFERED)); + } + + // ============================================================ + // Producer harness + // ============================================================ + + private interface Producer { + void run(int threadId) throws Exception; + } + + /** Runs one producer per thread, released simultaneously, and rethrows any producer failure. */ + private static void runProducers(int threads, Producer producer) throws Exception { + ExecutorService executor = Executors.newFixedThreadPool(threads); + try { + CountDownLatch start = new CountDownLatch(1); + List> futures = new ArrayList<>(); + for (int t = 0; t < threads; t++) { + int threadId = t; + futures.add( + executor.submit( + () -> { + start.await(); + producer.run(threadId); + return null; + })); + } + start.countDown(); + for (Future future : futures) { + future.get(30, TimeUnit.SECONDS); + } + } finally { + executor.shutdown(); + } + } +} From 34b84cb7061061945d7aa1306cd8931de0df57a3 Mon Sep 17 00:00:00 2001 From: Hazel Date: Wed, 15 Jul 2026 17:06:21 -0700 Subject: [PATCH 40/65] Add BillingS3ExportIntegrationTest --- .../v1/BillingS3ExportIntegrationTest.java | 217 ++++++++++++++++++ .../testresource/S3MockTestResource.java | 76 ++++++ 2 files changed, 293 insertions(+) create mode 100644 src/test/java/io/stargate/sgv2/jsonapi/api/v1/BillingS3ExportIntegrationTest.java create mode 100644 src/test/java/io/stargate/sgv2/jsonapi/testresource/S3MockTestResource.java diff --git a/src/test/java/io/stargate/sgv2/jsonapi/api/v1/BillingS3ExportIntegrationTest.java b/src/test/java/io/stargate/sgv2/jsonapi/api/v1/BillingS3ExportIntegrationTest.java new file mode 100644 index 0000000000..c32529fcc6 --- /dev/null +++ b/src/test/java/io/stargate/sgv2/jsonapi/api/v1/BillingS3ExportIntegrationTest.java @@ -0,0 +1,217 @@ +package io.stargate.sgv2.jsonapi.api.v1; + +import static io.restassured.RestAssured.given; +import static io.stargate.sgv2.jsonapi.api.v1.ResponseAssertions.responseIsDDLSuccess; +import static io.stargate.sgv2.jsonapi.api.v1.ResponseAssertions.responseIsWriteSuccess; +import static org.assertj.core.api.Assertions.assertThat; +import static org.awaitility.Awaitility.await; +import static org.hamcrest.Matchers.is; + +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import io.quarkus.test.common.WithTestResource; +import io.quarkus.test.junit.QuarkusIntegrationTest; +import io.stargate.sgv2.jsonapi.testresource.DseTestResource; +import io.stargate.sgv2.jsonapi.testresource.S3MockTestResource; +import java.net.URI; +import java.time.Duration; +import java.util.ArrayList; +import java.util.HashSet; +import java.util.List; +import java.util.Set; +import java.util.regex.Pattern; +import org.junit.jupiter.api.Test; +import software.amazon.awssdk.auth.credentials.AwsBasicCredentials; +import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider; +import software.amazon.awssdk.regions.Region; +import software.amazon.awssdk.services.s3.S3Client; +import software.amazon.awssdk.services.s3.model.S3Object; + +/** + * End-to-end test of the billing S3 export: real vectorize commands (via {@code + * CustomITEmbeddingProvider}) emit {@code billing.events} lines, and the installed {@code + * BillingS3LogHandler} must land them in the S3Mock bucket as time-partitioned NDJSON objects. + * + *

{@link S3MockTestResource} enables the export with small thresholds (count seal 5, age sweep + * 2s) and turns on the {@code billing-events-logging} feature flag. + */ +@QuarkusIntegrationTest +@WithTestResource(value = DseTestResource.class) +@WithTestResource(value = S3MockTestResource.class) +public class BillingS3ExportIntegrationTest extends AbstractKeyspaceIntegrationTestBase { + + private static final String COLLECTION = "billing_export_collection"; + + /** Every vectorize call emits at least one billing event, so lines >= documents. */ + private static final int DOCUMENTS = 10; + + private static final Pattern KEY_PATTERN = + Pattern.compile("billing-events/\\d{4}/\\d{2}/\\d{2}/\\d{2}/\\d{2}/[0-9a-f-]{36}\\.jsonl"); + + /** Wire contract of {@code BillingEventType}: billing consumers key on these exact values. */ + private static final Set EVENT_TYPES = + Set.of( + "internal_model_total_tokens", + "external_model_total_tokens", + "internal_model_egress_bytes", + "external_model_egress_bytes", + "internal_model_ingress_bytes", + "external_model_ingress_bytes"); + + private static final ObjectMapper MAPPER = new ObjectMapper(); + + @Test + public void billingEventsLandInS3AsNdjson() throws Exception { + createVectorizeCollection(); + for (int i = 0; i < DOCUMENTS; i++) { + insertDocumentWithVectorize(i); + } + + try (S3Client s3 = verificationClient()) { + // The count seal ships full batches immediately; the 2s age tick sweeps the remainder. + await() + .atMost(Duration.ofSeconds(60)) + .pollInterval(Duration.ofSeconds(2)) + .untilAsserted( + () -> assertThat(exportedLines(s3)).hasSizeGreaterThanOrEqualTo(DOCUMENTS)); + + // Object layout: time-partitioned keys and NDJSON content type. + List objects = exportObjects(s3); + assertThat(objects).isNotEmpty(); + for (S3Object object : objects) { + assertThat(object.key()).matches(KEY_PATTERN); + } + var head = s3.headObject(b -> b.bucket(S3MockTestResource.BUCKET).key(objects.get(0).key())); + assertThat(head.contentType()).isEqualTo("application/x-ndjson"); + + // Every line is a self-contained billing event with the expected shape; ids never repeat + // across objects. (region/resource_id may be absent locally and are not asserted.) + List lines = exportedLines(s3); + Set seenIds = new HashSet<>(); + for (String line : lines) { + JsonNode event = MAPPER.readTree(line); + String id = event.path("id").asText(); + assertThat(id).isNotBlank(); + assertThat(seenIds.add(id)) + .as("billing event id duplicated across export: %s", id) + .isTrue(); + assertThat(event.path("timestamp").asText()).isNotBlank(); + assertThat(event.path("product").asText()).isEqualTo("serverless"); + assertThat(event.path("event_type").asText()).isIn(EVENT_TYPES); + JsonNode properties = event.path("properties"); + assertThat(properties.path("usage").isIntegralNumber()).isTrue(); + assertThat(properties.path("usage").asLong()).isGreaterThanOrEqualTo(0L); + assertThat(properties.path("resource_type").asText()).isEqualTo("serverless_database"); + assertThat(properties.path("provider").asText()).isEqualTo("custom"); + // The billed model is what the provider reports in ModelUsage — for the IT provider that + // is its internal model config ("test-model"), not the createCollection modelName. + assertThat(properties.path("model").asText()).isEqualTo("test-model"); + } + } + + // The delivery counters on /metrics must agree that the export is alive. + await() + .atMost(Duration.ofSeconds(10)) + .untilAsserted( + () -> { + String metrics = + given().when().get("/metrics").then().statusCode(200).extract().asString(); + double flushedTotal = + metrics + .lines() + .filter(line -> line.startsWith("billing_s3_events_flushed_total")) + .mapToDouble( + line -> Double.parseDouble(line.substring(line.lastIndexOf(' ') + 1))) + .sum(); + assertThat(flushedTotal).isGreaterThanOrEqualTo(DOCUMENTS); + }); + } + + // ============================================================ + // Command helpers + // ============================================================ + + private void createVectorizeCollection() { + givenHeadersPostJsonThenOk( + """ + { + "createCollection": { + "name": "%s", + "options": { + "vector": { + "metric": "cosine", + "dimension": 5, + "service": { + "provider": "custom", + "modelName": "text-embedding-ada-002", + "authentication": { + "providerKey" : "shared_creds.providerKey" + }, + "parameters": { + "projectId": "test project" + } + } + } + } + } + } + """ + .formatted(COLLECTION)) + .body("$", responseIsDDLSuccess()) + .body("status.ok", is(1)); + } + + private void insertDocumentWithVectorize(int i) { + String json = + """ + { + "insertOne": { + "document": { + "_id": "doc-%d", + "description": "billing export test document %d", + "$vectorize": "billing export test document %d" + } + } + } + """ + .formatted(i, i, i); + givenHeadersAndJson(json) + .when() + .post(CollectionResource.BASE_PATH, keyspaceName, COLLECTION) + .then() + .statusCode(200) + .body("$", responseIsWriteSuccess()); + } + + // ============================================================ + // S3 verification helpers + // ============================================================ + + private static S3Client verificationClient() { + return S3Client.builder() + .region(Region.of(S3MockTestResource.BUCKET_REGION)) + .credentialsProvider( + StaticCredentialsProvider.create( + AwsBasicCredentials.create( + S3MockTestResource.ACCESS_KEY, S3MockTestResource.SECRET_KEY))) + .endpointOverride(URI.create(S3MockTestResource.endpoint())) + .forcePathStyle(true) + .build(); + } + + private static List exportObjects(S3Client s3) { + return s3.listObjectsV2(b -> b.bucket(S3MockTestResource.BUCKET).prefix("billing-events/")) + .contents(); + } + + private static List exportedLines(S3Client s3) { + List lines = new ArrayList<>(); + for (S3Object object : exportObjects(s3)) { + String body = + s3.getObjectAsBytes(b -> b.bucket(S3MockTestResource.BUCKET).key(object.key())) + .asUtf8String(); + body.lines().filter(line -> !line.isBlank()).forEach(lines::add); + } + return lines; + } +} diff --git a/src/test/java/io/stargate/sgv2/jsonapi/testresource/S3MockTestResource.java b/src/test/java/io/stargate/sgv2/jsonapi/testresource/S3MockTestResource.java new file mode 100644 index 0000000000..533f1d640f --- /dev/null +++ b/src/test/java/io/stargate/sgv2/jsonapi/testresource/S3MockTestResource.java @@ -0,0 +1,76 @@ +package io.stargate.sgv2.jsonapi.testresource; + +import com.adobe.testing.s3mock.testcontainers.S3MockContainer; +import io.quarkus.test.common.QuarkusTestResourceLifecycleManager; +import java.util.HashMap; +import java.util.Map; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Starts an S3Mock container and enables the billing S3 export against it, with small batch + * thresholds so tests see objects quickly. Used by {@code BillingS3ExportIntegrationTest} alongside + * {@link DseTestResource}. + * + *

The returned properties reach the application under test (a separate process for + * {@code @QuarkusIntegrationTest}); mirroring them as system properties follows the {@link + * StargateTestResource} pattern so the whole test environment sees the same values. + */ +public class S3MockTestResource implements QuarkusTestResourceLifecycleManager { + + private static final Logger LOG = LoggerFactory.getLogger(S3MockTestResource.class); + + /** Container tag; keep in sync with the {@code s3mock-testcontainers} version in pom.xml. */ + private static final String S3MOCK_VERSION = "5.1.0"; + + public static final String BUCKET = "billing-events-it"; + public static final String BUCKET_REGION = "us-east-1"; + public static final String ACCESS_KEY = "s3mock-test"; + public static final String SECRET_KEY = "s3mock-test"; + + private static volatile String httpEndpoint; + + private S3MockContainer container; + + /** HTTP endpoint of the running S3Mock, for the test-side verification client. */ + public static String endpoint() { + if (httpEndpoint == null) { + throw new IllegalStateException("S3MockTestResource has not been started"); + } + return httpEndpoint; + } + + @Override + public Map start() { + container = new S3MockContainer(S3MOCK_VERSION).withInitialBuckets(BUCKET); + container.start(); + httpEndpoint = container.getHttpEndpoint(); + + Map props = new HashMap<>(); + props.put("stargate.jsonapi.billing.s3.enabled", "true"); + props.put("stargate.jsonapi.billing.s3.bucket", BUCKET); + props.put("stargate.jsonapi.billing.s3.bucket-region", BUCKET_REGION); + props.put("stargate.jsonapi.billing.s3.endpoint-override", httpEndpoint); + // Small thresholds so the export flushes promptly: count seal at 5, age sweep every 2s. + props.put("stargate.jsonapi.billing.s3.max-events", "5"); + props.put("stargate.jsonapi.billing.s3.max-age", "PT2S"); + props.put("stargate.jsonapi.billing.s3.shutdown-timeout", "PT5S"); + // The producer side (DefaultBilling) is feature-flagged off by default. + props.put("stargate.feature.flags.billing-events-logging", "true"); + // The uploader resolves credentials from the SDK default chain, whose first stop is the + // system-property provider. S3Mock accepts any signed request. + props.put("aws.accessKeyId", ACCESS_KEY); + props.put("aws.secretAccessKey", SECRET_KEY); + + props.forEach(System::setProperty); + LOG.info("S3Mock started for billing export IT: endpoint={}, bucket={}", httpEndpoint, BUCKET); + return props; + } + + @Override + public void stop() { + if (container != null) { + container.stop(); + } + } +} From 994bd1c510870d1ecac5c86631f80caf79a3a97e Mon Sep 17 00:00:00 2001 From: Hazel Date: Wed, 15 Jul 2026 17:35:20 -0700 Subject: [PATCH 41/65] Add java doc for S3BatchUploader --- .../sgv2/jsonapi/service/provider/S3BatchUploader.java | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/S3BatchUploader.java b/src/main/java/io/stargate/sgv2/jsonapi/service/provider/S3BatchUploader.java index 2043a7bc5b..b617c92176 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/S3BatchUploader.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/service/provider/S3BatchUploader.java @@ -1,5 +1,6 @@ package io.stargate.sgv2.jsonapi.service.provider; +import com.google.common.annotations.VisibleForTesting; import io.smallrye.mutiny.Uni; import java.net.URI; import java.nio.charset.StandardCharsets; @@ -76,6 +77,10 @@ public static S3BatchUploader create( public Uni upload(BillingQueue.Batch batch) { String key = objectKey(batch.oldestEventAt(), UUID.randomUUID()); byte[] body = toNdjson(batch.lines()); + // No .retry() here: unconfigured, S3AsyncClient already retries (default LegacyRetryStrategy — + // 3 retries / 4 attempts). See + // https://docs.aws.amazon.com/sdk-for-java/latest/developer-guide/retry-strategy.html + // and see https://github.com/aws/aws-sdk-java-v2/issues/6987 for future change. return Uni.createFrom() .completionStage( () -> @@ -90,10 +95,12 @@ public Uni upload(BillingQueue.Batch batch) { .thenAccept(resp -> {})); } + @VisibleForTesting static String objectKey(Instant timestamp, UUID id) { return PATH_PREFIX + "/" + KEY_TIME_FORMAT.format(timestamp) + "/" + id + ".jsonl"; } + @VisibleForTesting static byte[] toNdjson(List lines) { StringBuilder sb = new StringBuilder(); for (String line : lines) { From 4748005251452a780fc2d2213661654ec6dd37c6 Mon Sep 17 00:00:00 2001 From: Hazel Date: Wed, 15 Jul 2026 18:36:08 -0700 Subject: [PATCH 42/65] Add java comments for tests --- .../service/provider/BillingS3LogHandlerTest.java | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/src/test/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandlerTest.java b/src/test/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandlerTest.java index 6e43aff18c..92996643b3 100644 --- a/src/test/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandlerTest.java +++ b/src/test/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandlerTest.java @@ -34,9 +34,6 @@ * on close), the upload-concurrency gate, failure containment, and the at-most-once accounting * invariants under concurrent publish. The uploader is a programmable in-memory fake; real S3 I/O * is covered by {@code BillingS3ExportIntegrationTest}. - * - *

Ordering is only asserted in single-producer tests (the buffer is FIFO for a single thread); - * concurrency tests assert set-equality and counter reconciliation, never interleaving order. */ class BillingS3LogHandlerTest { @@ -228,6 +225,7 @@ void sealsByCountAndShipsExactBatch() { // record is older on purpose, so oldestEventAt must be the min, not the head). handler.publish(record(T0.plusSeconds(5), "{\"e\":1}")); handler.publish(record(T0, "{\"e\":2}")); + // Asserts the condition holds for the whole window — i.e. that an async flush did NOT happen await() .during(Duration.ofMillis(200)) .atMost(Duration.ofSeconds(2)) @@ -297,6 +295,7 @@ void noShipmentBelowSealUntilAgeTick() { void ageTickIsScheduledForReal() { var uploader = new RecordingUploader(); var registry = new SimpleMeterRegistry(); + // Raw constructor: newHandler() pins maxAge to NEVER; this is the one test that wants a live tick. var handler = new BillingS3LogHandler( uploader, @@ -311,8 +310,8 @@ void ageTickIsScheduledForReal() { handler.publish(record("{\"e\":1}")); // No seal is reached; only the scheduled fixed-rate tick can ship this line. - await().atMost(AWAIT).untilAsserted(() -> assertThat(uploader.batches).isNotEmpty()); - assertThat(uploader.allLines()).containsExactly("{\"e\":1}"); + await().atMost(AWAIT).untilAsserted(() -> assertThat(uploader.batches).hasSize(1)); + assertThat(uploader.batches.get(0).lines()).containsExactly("{\"e\":1}"); } finally { handler.close(); } From 7e2408f8f3bba5f0f506e5d21520774c056a8958 Mon Sep 17 00:00:00 2001 From: Hazel Date: Wed, 15 Jul 2026 20:49:29 -0700 Subject: [PATCH 43/65] Add test --- .../provider/BillingS3LogHandlerTest.java | 34 +++++++++---------- 1 file changed, 16 insertions(+), 18 deletions(-) diff --git a/src/test/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandlerTest.java b/src/test/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandlerTest.java index 92996643b3..9d83cbeba8 100644 --- a/src/test/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandlerTest.java +++ b/src/test/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandlerTest.java @@ -295,7 +295,8 @@ void noShipmentBelowSealUntilAgeTick() { void ageTickIsScheduledForReal() { var uploader = new RecordingUploader(); var registry = new SimpleMeterRegistry(); - // Raw constructor: newHandler() pins maxAge to NEVER; this is the one test that wants a live tick. + // Raw constructor: newHandler() pins maxAge to NEVER; this is the one test that wants a live + // tick. var handler = new BillingS3LogHandler( uploader, @@ -409,28 +410,25 @@ void closeDrainsRemainderAndClosesUploader() { } @Test - void closeTimeoutCountsAbandoned() throws Exception { + void closeTimeoutCountsAbandoned() { var uploader = new RecordingUploader(); var registry = new SimpleMeterRegistry(); var handler = newHandler(uploader, registry, 1, 1_000_000, 10, 1, Duration.ofMillis(200)); - try { - uploader.mode = RecordingUploader.Mode.HOLD; - handler.publish(record("{\"e\":1}")); - // Wait until the first batch is in flight (and held) so the queued remainder is exact. - await().atMost(AWAIT).untilAsserted(() -> assertThat(uploader.batches).hasSize(1)); - handler.publish(record("{\"e\":2}")); - handler.publish(record("{\"e\":3}")); - long startNanos = System.nanoTime(); - handler.close(); - Duration elapsed = Duration.ofNanos(System.nanoTime() - startNanos); + uploader.mode = RecordingUploader.Mode.HOLD; + handler.publish(record("{\"e\":1}")); + // Wait until the first batch is in flight (and held) so the queued remainder is exact. + await().atMost(AWAIT).untilAsserted(() -> assertThat(uploader.batches).hasSize(1)); + handler.publish(record("{\"e\":2}")); + handler.publish(record("{\"e\":3}")); - assertThat(elapsed).isLessThan(Duration.ofSeconds(3)); - assertThat(counter(registry, DROPPED, "reason", "shutdown")).isEqualTo(2.0); - assertThat(uploader.closed).isTrue(); - } finally { - uploader.releaseAllAndComplete(); - } + long startNanos = System.nanoTime(); + handler.close(); + Duration elapsed = Duration.ofNanos(System.nanoTime() - startNanos); + + assertThat(elapsed).isLessThan(Duration.ofSeconds(3)); + assertThat(counter(registry, DROPPED, "reason", "shutdown")).isEqualTo(2.0); + assertThat(uploader.closed).isTrue(); } @Test From 62ab29cb1b7342b32906638387fdaa65f79b033c Mon Sep 17 00:00:00 2001 From: Hazel Date: Wed, 15 Jul 2026 21:49:08 -0700 Subject: [PATCH 44/65] update test --- .../jsonapi/service/provider/BillingS3LogHandlerTest.java | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/test/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandlerTest.java b/src/test/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandlerTest.java index 9d83cbeba8..bf5b8fb168 100644 --- a/src/test/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandlerTest.java +++ b/src/test/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandlerTest.java @@ -403,7 +403,8 @@ void closeDrainsRemainderAndClosesUploader() { handler.close(); // close() is synchronous: by the time it returns the drain has settled and counted. - assertThat(uploader.allLines()).containsExactly("{\"e\":1}", "{\"e\":2}", "{\"e\":3}"); + assertThat(uploader.allLines()) + .containsExactlyInAnyOrder("{\"e\":1}", "{\"e\":2}", "{\"e\":3}"); assertThat(uploader.closed).isTrue(); assertThat(counter(registry, FLUSHED)).isEqualTo(3.0); assertThat(counter(registry, DROPPED, "reason", "shutdown")).isZero(); @@ -517,7 +518,7 @@ void overflowAccountingReconciles() throws Exception { } @Test - void publishNeverBlocksWhenUploaderStalls() throws Exception { + void publishNeverBlocksWhenUploaderStalls() { var uploader = new RecordingUploader(); var registry = new SimpleMeterRegistry(); var handler = newHandler(uploader, registry, 1, 1_000_000, 8, 1); From a751ddec22bed9ac119822a6dc2a6bf56f82f330 Mon Sep 17 00:00:00 2001 From: Hazel Date: Wed, 15 Jul 2026 22:05:22 -0700 Subject: [PATCH 45/65] update test order --- .../provider/BillingS3LogHandlerTest.java | 140 +++++++++--------- 1 file changed, 72 insertions(+), 68 deletions(-) diff --git a/src/test/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandlerTest.java b/src/test/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandlerTest.java index bf5b8fb168..e805894a5d 100644 --- a/src/test/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandlerTest.java +++ b/src/test/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandlerTest.java @@ -447,76 +447,9 @@ void closeIsIdempotentAndPublishAfterCloseIsSafe() { } // ============================================================ - // Concurrency — invariant style, no interleaving assertions + // Async pipeline — gate, chain liveness, back-pressure (deterministic, single driver thread) // ============================================================ - @Test - void multiProducerNoLossNoDuplication() throws Exception { - var uploader = new RecordingUploader(); - var registry = new SimpleMeterRegistry(); - var handler = newHandler(uploader, registry, 50, 1_000_000_000L, 10_000, 4); - - int threads = 8; - int perThread = 500; - Set published = ConcurrentHashMap.newKeySet(); - runProducers( - threads, - (threadId) -> { - for (int i = 0; i < perThread; i++) { - String line = "{\"t\":" + threadId + ",\"i\":" + i + "}"; - published.add(line); - handler.publish(record(line)); - } - }); - handler.close(); - - List delivered = uploader.allLines(); - assertThat(delivered).hasSize(threads * perThread); - assertThat(new HashSet<>(delivered)).isEqualTo(published); - assertThat(counter(registry, OFFERED)).isEqualTo(threads * perThread); - assertThat(counter(registry, FLUSHED)).isEqualTo(threads * perThread); - assertThat(counter(registry, DROPPED, "reason", "capacity")).isZero(); - assertThat(counter(registry, DROPPED, "reason", "shutdown")).isZero(); - } - - @Test - void overflowAccountingReconciles() throws Exception { - var uploader = new RecordingUploader(); - var registry = new SimpleMeterRegistry(); - // No seal is ever reached (count seal above capacity, byte seal huge): nothing drains while - // producers run, so every line beyond the 64-slot buffer is a deterministic capacity drop. - var handler = newHandler(uploader, registry, 1000, 1_000_000_000L, 64, 4); - - int threads = 4; - int perThread = 500; - Set published = ConcurrentHashMap.newKeySet(); - runProducers( - threads, - (threadId) -> { - for (int i = 0; i < perThread; i++) { - String line = "{\"t\":" + threadId + ",\"i\":" + i + "}"; - published.add(line); - handler.publish(record(line)); - } - }); - - uploader.releaseAllAndComplete(); - handler.close(); - - double offered = counter(registry, OFFERED); - double flushed = counter(registry, FLUSHED); - double droppedCapacity = counter(registry, DROPPED, "reason", "capacity"); - double droppedShutdown = counter(registry, DROPPED, "reason", "shutdown"); - assertThat(offered).isEqualTo(threads * perThread); - assertThat(flushed).isEqualTo(64.0); - assertThat(droppedCapacity).isEqualTo(threads * perThread - 64.0); - assertThat(flushed + droppedCapacity + droppedShutdown).isEqualTo(offered); - - List delivered = uploader.allLines(); - assertThat(delivered).doesNotHaveDuplicates(); - assertThat(published).containsAll(delivered); - } - @Test void publishNeverBlocksWhenUploaderStalls() { var uploader = new RecordingUploader(); @@ -612,6 +545,77 @@ void settledUploadChainsNextBatchWithoutNewPublish() throws Exception { } } + // ============================================================ + // Concurrent producers — accounting invariants under racing publish (interleaving-agnostic) + // ============================================================ + + @Test + void multiProducerNoLossNoDuplication() throws Exception { + var uploader = new RecordingUploader(); + var registry = new SimpleMeterRegistry(); + var handler = newHandler(uploader, registry, 50, 1_000_000_000L, 10_000, 4); + + int threads = 8; + int perThread = 500; + Set published = ConcurrentHashMap.newKeySet(); + runProducers( + threads, + (threadId) -> { + for (int i = 0; i < perThread; i++) { + String line = "{\"t\":" + threadId + ",\"i\":" + i + "}"; + published.add(line); + handler.publish(record(line)); + } + }); + handler.close(); + + List delivered = uploader.allLines(); + assertThat(delivered).hasSize(threads * perThread); + assertThat(new HashSet<>(delivered)).isEqualTo(published); + assertThat(counter(registry, OFFERED)).isEqualTo(threads * perThread); + assertThat(counter(registry, FLUSHED)).isEqualTo(threads * perThread); + assertThat(counter(registry, DROPPED, "reason", "capacity")).isZero(); + assertThat(counter(registry, DROPPED, "reason", "shutdown")).isZero(); + } + + @Test + void overflowAccountingReconciles() throws Exception { + var uploader = new RecordingUploader(); + var registry = new SimpleMeterRegistry(); + // No seal is ever reached (count seal above capacity, byte seal huge): nothing drains while + // producers run, so every line beyond the 64-slot buffer is a deterministic capacity drop. + var handler = newHandler(uploader, registry, 1000, 1_000_000_000L, 64, 4); + + int threads = 4; + int perThread = 500; + Set published = ConcurrentHashMap.newKeySet(); + runProducers( + threads, + (threadId) -> { + for (int i = 0; i < perThread; i++) { + String line = "{\"t\":" + threadId + ",\"i\":" + i + "}"; + published.add(line); + handler.publish(record(line)); + } + }); + + uploader.releaseAllAndComplete(); + handler.close(); + + double offered = counter(registry, OFFERED); + double flushed = counter(registry, FLUSHED); + double droppedCapacity = counter(registry, DROPPED, "reason", "capacity"); + double droppedShutdown = counter(registry, DROPPED, "reason", "shutdown"); + assertThat(offered).isEqualTo(threads * perThread); + assertThat(flushed).isEqualTo(64.0); + assertThat(droppedCapacity).isEqualTo(threads * perThread - 64.0); + assertThat(flushed + droppedCapacity + droppedShutdown).isEqualTo(offered); + + List delivered = uploader.allLines(); + assertThat(delivered).doesNotHaveDuplicates(); + assertThat(published).containsAll(delivered); + } + @Test void closeRacingProducersNeverHangsAndReconciles() throws Exception { var uploader = new RecordingUploader(); From 78265781f9689d406538e9db735c6e97a20b6d53 Mon Sep 17 00:00:00 2001 From: Hazel Date: Thu, 16 Jul 2026 07:40:27 -0700 Subject: [PATCH 46/65] add comments --- .../jsonapi/service/provider/BillingS3HandlerInstaller.java | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3HandlerInstaller.java b/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3HandlerInstaller.java index c1057551e3..da0086a72d 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3HandlerInstaller.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3HandlerInstaller.java @@ -70,6 +70,8 @@ void onStop(@Observes ShutdownEvent event) { return; } Logger.getLogger(BILLING_LOGGER_NAME).removeHandler(this.handler); + // close() isn't expected to throw, but if it does (e.g. client.close() failing), letting it + // propagate would disrupt other components' cleanup in Quarkus's shutdown sequence. try { this.handler.close(); } catch (Exception e) { From 7844eccf02089332a75c4a6a69c8dd539e2e60c2 Mon Sep 17 00:00:00 2001 From: Hazel Date: Thu, 16 Jul 2026 08:06:27 -0700 Subject: [PATCH 47/65] update tests --- .../service/provider/BillingS3LogHandlerTest.java | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/src/test/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandlerTest.java b/src/test/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandlerTest.java index e805894a5d..1911ee7817 100644 --- a/src/test/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandlerTest.java +++ b/src/test/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandlerTest.java @@ -493,15 +493,13 @@ void concurrencyGateCapsParallelUploads() throws Exception { await().atMost(AWAIT).untilAsserted(() -> assertThat(uploader.inFlight.get()).isEqualTo(2)); // Each release lets exactly the next batch through, one at a time. - while (uploader.batches.size() < 10) { + for (int expected = 3; expected <= 10; expected++) { uploader.releaseOne(); - int expected = Math.min(uploader.batches.size() + 1, 10); - await() - .atMost(AWAIT) - .untilAsserted( - () -> assertThat(uploader.batches.size()).isGreaterThanOrEqualTo(expected)); + int size = expected; // fresh effectively-final binding for the lambda + await().atMost(AWAIT).untilAsserted(() -> assertThat(uploader.batches).hasSize(size)); } + // maxInFlight = RecordingUploader's high-water mark: uploads peaked at the gate cap of 2. assertThat(uploader.maxInFlight.get()).isEqualTo(2); assertThat(uploader.allLines()).hasSize(10).doesNotHaveDuplicates(); } finally { From 650d8a8904c6a38b1c36a7ce44686b2dc49b568e Mon Sep 17 00:00:00 2001 From: Hazel Date: Thu, 16 Jul 2026 08:25:17 -0700 Subject: [PATCH 48/65] update test --- .../jsonapi/service/provider/BillingS3LogHandlerTest.java | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/test/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandlerTest.java b/src/test/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandlerTest.java index 1911ee7817..2a11deed3e 100644 --- a/src/test/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandlerTest.java +++ b/src/test/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandlerTest.java @@ -520,6 +520,8 @@ void settledUploadChainsNextBatchWithoutNewPublish() throws Exception { } // Single slot: exactly one upload starts, the two other sealed batches wait behind it. + // First await = arrival: the one upload has started. Second = the gate holds: batches stays + // at exactly one for 200ms, proving concurrency=1 keeps the other two sealed batches queued. await().atMost(AWAIT).untilAsserted(() -> assertThat(uploader.batches).hasSize(1)); await() .during(Duration.ofMillis(200)) @@ -536,7 +538,9 @@ void settledUploadChainsNextBatchWithoutNewPublish() throws Exception { await().atMost(AWAIT).untilAsserted(() -> assertThat(uploader.batches).hasSize(3)); assertThat(uploader.batches.get(2).lines()).containsExactly("{\"i\":5}", "{\"i\":6}"); + // Settling the last in-flight upload with an empty buffer chains nothing further. uploader.releaseOne(); + assertThat(uploader.batches).hasSize(3); } finally { uploader.releaseAllAndComplete(); handler.close(); From de68c32e9e3cf886e7862fb657da5b0774d48dc0 Mon Sep 17 00:00:00 2001 From: Hazel Date: Thu, 16 Jul 2026 12:39:15 -0700 Subject: [PATCH 49/65] update test --- .../provider/BillingS3LogHandlerTest.java | 48 +++++++++++++++---- 1 file changed, 40 insertions(+), 8 deletions(-) diff --git a/src/test/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandlerTest.java b/src/test/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandlerTest.java index 2a11deed3e..e4b085094f 100644 --- a/src/test/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandlerTest.java +++ b/src/test/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandlerTest.java @@ -28,6 +28,8 @@ import java.util.logging.LogRecord; import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.Test; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; /** * Unit tests for {@link BillingS3LogHandler}: the flush triggers (seal on publish, age tick, drain @@ -37,6 +39,8 @@ */ class BillingS3LogHandlerTest { + private static final Logger LOG = LoggerFactory.getLogger(BillingS3LogHandlerTest.class); + private static final Duration AWAIT = Duration.ofSeconds(10); private static final Duration NEVER = Duration.ofHours(1); private static final Instant T0 = Instant.parse("2026-05-20T14:23:11Z"); @@ -618,6 +622,19 @@ void overflowAccountingReconciles() throws Exception { assertThat(published).containsAll(delivered); } + /** + * {@code close()} runs concurrently with in-flight {@code publish()} calls in production — a pod + * shutdown doesn't wait for request threads to go quiet first. This drives both at once: 4 + * threads publish flat out while the main thread calls {@code close()} mid-stream, then keeps the + * producers running a bit longer so some publishes land after close() too. + * + *

Expected: no publish ever throws (the handler must stay safe under this race), close() + * returns instead of hanging, delivered lines are a duplicate-free subset of what was published, + * and the metrics reconcile as {@code flushed + dropped <= offered} rather than {@code ==}. The + * gap is expected, not a bug: a publish can land after close() takes its final buffer snapshot, + * so that line is neither delivered nor counted as dropped — see {@link BillingMetrics}'s class + * doc for this same at-most-once slippage. The log line below reports the exact gap each run. + */ @Test void closeRacingProducersNeverHangsAndReconciles() throws Exception { var uploader = new RecordingUploader(); @@ -625,8 +642,9 @@ void closeRacingProducersNeverHangsAndReconciles() throws Exception { var handler = newHandler(uploader, registry, 5, 1_000_000_000L, 1000, 4, Duration.ofSeconds(1)); int threads = 4; - // Producers run flat out until stopped; the cap only bounds memory and assertion cost, sized - // so publishing is still in full flight when close() lands ~50ms in. + // Producers normally exit on the stop flag below. The cap bounds the sad path (a hung close + // never reaches stop.set) so no producer spins forever — executor.shutdown() does not interrupt + // running tasks. Overlap with close() is guaranteed by the published.size() gate below. int perThreadCap = 200_000; Set published = ConcurrentHashMap.newKeySet(); List producerErrors = new CopyOnWriteArrayList<>(); @@ -651,7 +669,10 @@ void closeRacingProducersNeverHangsAndReconciles() throws Exception { })); } - Thread.sleep(50); // let producers overlap the close below + // Land close() deterministically amid in-flight publishes: wait until producers have flooded + // the pipeline (2x the 1000-slot buffer → buffer full, overflow dropping, uploads gated), not a + // wall-clock guess. AWAIT only bounds a stuck ramp-up. + await().atMost(AWAIT).until(() -> published.size() >= 2_000); handler.close(); stop.set(true); for (Future future : futures) { @@ -659,9 +680,10 @@ void closeRacingProducersNeverHangsAndReconciles() throws Exception { } executor.shutdown(); - // publish must never throw, close must return, and the books must stay consistent — lines - // published after close may be dropped without being counted, so the reconciliation is <=. - // Plain java.util.Set operations keep these checks O(n); AssertJ's containsAll would scan. + // publish must never throw, close must return. Post-close publishes can settle after close()'s + // final buffer snapshot, uncounted, so accounting reconciles with <=, not == (the log shows + // that gap). + // Plain Set ops keep these checks O(n); AssertJ's containsAll would scan, which is O(n^2). assertThat(producerErrors).isEmpty(); List delivered = uploader.allLines(); Set deliveredSet = new HashSet<>(delivered); @@ -669,11 +691,21 @@ void closeRacingProducersNeverHangsAndReconciles() throws Exception { assertThat(published.containsAll(deliveredSet)) .as("every delivered line must have been published") .isTrue(); + double offered = counter(registry, OFFERED); double flushed = counter(registry, FLUSHED); double droppedCapacity = counter(registry, DROPPED, "reason", "capacity"); double droppedShutdown = counter(registry, DROPPED, "reason", "shutdown"); - assertThat(flushed + droppedCapacity + droppedShutdown) - .isLessThanOrEqualTo(counter(registry, OFFERED)); + double accounted = flushed + droppedCapacity + droppedShutdown; + LOG.info( + "closeRacing reconcile: offered={} flushed={} droppedCapacity={} droppedShutdown={}" + + " accounted={} unaccounted={}", + (long) offered, + (long) flushed, + (long) droppedCapacity, + (long) droppedShutdown, + (long) accounted, + (long) (offered - accounted)); + assertThat(accounted).isLessThanOrEqualTo(offered); } // ============================================================ From 483c2c2e8b5928d67635f63702b27c2d7e885c0a Mon Sep 17 00:00:00 2001 From: Hazel Date: Thu, 16 Jul 2026 14:41:00 -0700 Subject: [PATCH 50/65] update S3MockTestResource --- .../jsonapi/testresource/S3MockTestResource.java | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/src/test/java/io/stargate/sgv2/jsonapi/testresource/S3MockTestResource.java b/src/test/java/io/stargate/sgv2/jsonapi/testresource/S3MockTestResource.java index 533f1d640f..9ef3c89e76 100644 --- a/src/test/java/io/stargate/sgv2/jsonapi/testresource/S3MockTestResource.java +++ b/src/test/java/io/stargate/sgv2/jsonapi/testresource/S3MockTestResource.java @@ -30,7 +30,7 @@ public class S3MockTestResource implements QuarkusTestResourceLifecycleManager { private static volatile String httpEndpoint; - private S3MockContainer container; + private static volatile S3MockContainer container; /** HTTP endpoint of the running S3Mock, for the test-side verification client. */ public static String endpoint() { @@ -40,6 +40,19 @@ public static String endpoint() { return httpEndpoint; } + /** + * Stops the S3Mock container, leaving nothing listening on the exported endpoint: every upload + * from then on fails with connection-refused, like an S3 outage. One-way for the whole test class + * (a restart would map a new port, unreachable through the app's fixed endpoint-override), so + * only the last test may call this. + */ + public static void stopContainer() { + if (container == null) { + throw new IllegalStateException("S3MockTestResource has not been started"); + } + container.stop(); + } + @Override public Map start() { container = new S3MockContainer(S3MOCK_VERSION).withInitialBuckets(BUCKET); From 64905443c5c195aa625cec91e99dad8b0caf066a Mon Sep 17 00:00:00 2001 From: Hazel Date: Thu, 16 Jul 2026 14:48:34 -0700 Subject: [PATCH 51/65] update BillingS3ExportIntegrationTest --- .../v1/BillingS3ExportIntegrationTest.java | 72 +++++++++++++++---- 1 file changed, 60 insertions(+), 12 deletions(-) diff --git a/src/test/java/io/stargate/sgv2/jsonapi/api/v1/BillingS3ExportIntegrationTest.java b/src/test/java/io/stargate/sgv2/jsonapi/api/v1/BillingS3ExportIntegrationTest.java index c32529fcc6..47cfe5ed13 100644 --- a/src/test/java/io/stargate/sgv2/jsonapi/api/v1/BillingS3ExportIntegrationTest.java +++ b/src/test/java/io/stargate/sgv2/jsonapi/api/v1/BillingS3ExportIntegrationTest.java @@ -20,7 +20,10 @@ import java.util.List; import java.util.Set; import java.util.regex.Pattern; +import org.junit.jupiter.api.MethodOrderer; +import org.junit.jupiter.api.Order; import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.TestMethodOrder; import software.amazon.awssdk.auth.credentials.AwsBasicCredentials; import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider; import software.amazon.awssdk.regions.Region; @@ -34,10 +37,14 @@ * *

{@link S3MockTestResource} enables the export with small thresholds (count seal 5, age sweep * 2s) and turns on the {@code billing-events-logging} feature flag. + * + *

Methods are ordered: the last test stops the S3Mock container to prove a failing export never + * affects the data API, which kills S3 for the rest of the class — nothing may run after it. */ @QuarkusIntegrationTest @WithTestResource(value = DseTestResource.class) @WithTestResource(value = S3MockTestResource.class) +@TestMethodOrder(MethodOrderer.OrderAnnotation.class) public class BillingS3ExportIntegrationTest extends AbstractKeyspaceIntegrationTestBase { private static final String COLLECTION = "billing_export_collection"; @@ -113,18 +120,49 @@ public void billingEventsLandInS3AsNdjson() throws Exception { await() .atMost(Duration.ofSeconds(10)) .untilAsserted( - () -> { - String metrics = - given().when().get("/metrics").then().statusCode(200).extract().asString(); - double flushedTotal = - metrics - .lines() - .filter(line -> line.startsWith("billing_s3_events_flushed_total")) - .mapToDouble( - line -> Double.parseDouble(line.substring(line.lastIndexOf(' ') + 1))) - .sum(); - assertThat(flushedTotal).isGreaterThanOrEqualTo(DOCUMENTS); - }); + () -> + assertThat(metricTotal("billing_s3_events_flushed_total")) + .isGreaterThanOrEqualTo(DOCUMENTS)); + } + + /** + * Billing is a side-channel: a failing S3 export must never affect the data API. Stopping the + * S3Mock container leaves the endpoint dead — every upload from here on fails with + * connection-refused, like an S3 outage — yet inserts must keep returning normal write successes, + * and the failures must be counted rather than silently swallowed. The handler's failure + * accounting in isolation is covered by {@code BillingS3LogHandlerTest}; this proves the property + * end-to-end in the packaged app. + * + *

Must run last ({@link S3MockTestResource#stopContainer()} is one-way): any test needing a + * live S3 goes before this one. Reuses the collection created by the happy-path test. + * + *

{@code Integer.MAX_VALUE}, not a small sentinel, and deliberately the only ordered method: + * {@code OrderAnnotation} gives an unannotated method the default order {@code Integer.MAX_VALUE + * / 2}, so any newly added test with no {@code @Order} still sorts before this one. Do NOT lower + * this value — anything below the default would let such a test run after S3 is dead. + */ + @Test + @Order(Integer.MAX_VALUE) + public void exportFailureDoesNotAffectTheApi() { + S3MockTestResource.stopContainer(); + + // Each insert emits billing events whose upload will fail — yet every insert must still + // return a normal write success, because publish() is fire-and-forget and never waits on S3. + for (int i = 0; i < DOCUMENTS; i++) { + insertDocumentWithVectorize(DOCUMENTS + i); + } + + // Failures are counted, not silently swallowed. Uploads settle as failed only after the SDK + // exhausts its retries, so poll for the counter to move. + await() + .atMost(Duration.ofSeconds(60)) + .pollInterval(Duration.ofSeconds(2)) + .untilAsserted( + () -> assertThat(metricTotal("billing_s3_batches_failed_total")).isGreaterThan(0.0)); + + // The API is still healthy after the export has been failing for a while: one more insert + // succeeds exactly like the first. + insertDocumentWithVectorize(2 * DOCUMENTS); } // ============================================================ @@ -214,4 +252,14 @@ private static List exportedLines(S3Client s3) { } return lines; } + + /** Sum of one counter across all tag combinations on {@code /metrics} (0 when absent). */ + private static double metricTotal(String metricName) { + String metrics = given().when().get("/metrics").then().statusCode(200).extract().asString(); + return metrics + .lines() + .filter(line -> line.startsWith(metricName)) + .mapToDouble(line -> Double.parseDouble(line.substring(line.lastIndexOf(' ') + 1))) + .sum(); + } } From 9ebff9940278029b5ea132dc659953f64b3a27da Mon Sep 17 00:00:00 2001 From: Hazel Date: Mon, 20 Jul 2026 12:33:59 -0700 Subject: [PATCH 52/65] update path prefix --- .../sgv2/jsonapi/service/provider/S3BatchUploader.java | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/S3BatchUploader.java b/src/main/java/io/stargate/sgv2/jsonapi/service/provider/S3BatchUploader.java index b617c92176..dd29b535b2 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/S3BatchUploader.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/service/provider/S3BatchUploader.java @@ -20,8 +20,8 @@ /** Uploads sealed billing batches to S3 as NDJSON objects under time-partitioned keys. */ public class S3BatchUploader implements BillingS3LogHandler.AsyncBatchUploader { - // S3 object-key consistent identifier; TBD - static final String PATH_PREFIX = "billing-events"; + // S3 object-key consistent identifier + static final String PATH_PREFIX = "data-api"; private static final String NDJSON_CONTENT_TYPE = "application/x-ndjson"; // object key format private static final DateTimeFormatter KEY_TIME_FORMAT = From 5f61a284faab416d2dec13bcc1208b90a233655c Mon Sep 17 00:00:00 2001 From: Hazel Date: Mon, 20 Jul 2026 12:41:48 -0700 Subject: [PATCH 53/65] update path prefix test --- .../sgv2/jsonapi/service/provider/S3BatchUploaderTest.java | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/test/java/io/stargate/sgv2/jsonapi/service/provider/S3BatchUploaderTest.java b/src/test/java/io/stargate/sgv2/jsonapi/service/provider/S3BatchUploaderTest.java index 90e8405354..bf238a6f29 100644 --- a/src/test/java/io/stargate/sgv2/jsonapi/service/provider/S3BatchUploaderTest.java +++ b/src/test/java/io/stargate/sgv2/jsonapi/service/provider/S3BatchUploaderTest.java @@ -32,7 +32,7 @@ class S3BatchUploaderTest { private static final Duration AWAIT = Duration.ofSeconds(5); private static final Pattern KEY_PATTERN = - Pattern.compile("billing-events/2026/05/20/14/23/[0-9a-f-]{36}\\.jsonl"); + Pattern.compile("data-api/2026/05/20/14/23/[0-9a-f-]{36}\\.jsonl"); private static final String LINE_A = "{\"a\":1}"; private static final String LINE_B = "{\"b\":2}"; private static final BillingQueue.Batch BATCH = @@ -55,7 +55,7 @@ void objectKeyUsesPathPrefixAndUtcMinutePathFromTimestamp() { var id = UUID.fromString("8c0e9b8a-1d3a-4f6b-9c0d-1234567890ab"); var key = S3BatchUploader.objectKey(Instant.parse("2026-05-20T14:23:11.482Z"), id); assertThat(key) - .isEqualTo("billing-events/2026/05/20/14/23/8c0e9b8a-1d3a-4f6b-9c0d-1234567890ab.jsonl"); + .isEqualTo("data-api/2026/05/20/14/23/8c0e9b8a-1d3a-4f6b-9c0d-1234567890ab.jsonl"); } @Test From dc3c4963ae4a902158bab6e05674fcbbfb8475a7 Mon Sep 17 00:00:00 2001 From: Hazel Date: Mon, 20 Jul 2026 12:57:03 -0700 Subject: [PATCH 54/65] update path prefix test --- .../sgv2/jsonapi/api/v1/BillingS3ExportIntegrationTest.java | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/test/java/io/stargate/sgv2/jsonapi/api/v1/BillingS3ExportIntegrationTest.java b/src/test/java/io/stargate/sgv2/jsonapi/api/v1/BillingS3ExportIntegrationTest.java index 47cfe5ed13..b121a50540 100644 --- a/src/test/java/io/stargate/sgv2/jsonapi/api/v1/BillingS3ExportIntegrationTest.java +++ b/src/test/java/io/stargate/sgv2/jsonapi/api/v1/BillingS3ExportIntegrationTest.java @@ -53,7 +53,7 @@ public class BillingS3ExportIntegrationTest extends AbstractKeyspaceIntegrationT private static final int DOCUMENTS = 10; private static final Pattern KEY_PATTERN = - Pattern.compile("billing-events/\\d{4}/\\d{2}/\\d{2}/\\d{2}/\\d{2}/[0-9a-f-]{36}\\.jsonl"); + Pattern.compile("data-api/\\d{4}/\\d{2}/\\d{2}/\\d{2}/\\d{2}/[0-9a-f-]{36}\\.jsonl"); /** Wire contract of {@code BillingEventType}: billing consumers key on these exact values. */ private static final Set EVENT_TYPES = @@ -238,7 +238,7 @@ private static S3Client verificationClient() { } private static List exportObjects(S3Client s3) { - return s3.listObjectsV2(b -> b.bucket(S3MockTestResource.BUCKET).prefix("billing-events/")) + return s3.listObjectsV2(b -> b.bucket(S3MockTestResource.BUCKET).prefix("data-api/")) .contents(); } From ab0782306c6fcecede77eaa9390c448e60bbf08b Mon Sep 17 00:00:00 2001 From: Hazel Date: Tue, 21 Jul 2026 21:42:50 -0700 Subject: [PATCH 55/65] empty commit --- .../sgv2/jsonapi/api/v1/BillingS3ExportIntegrationTest.java | 1 + 1 file changed, 1 insertion(+) diff --git a/src/test/java/io/stargate/sgv2/jsonapi/api/v1/BillingS3ExportIntegrationTest.java b/src/test/java/io/stargate/sgv2/jsonapi/api/v1/BillingS3ExportIntegrationTest.java index b121a50540..6e4d835f0b 100644 --- a/src/test/java/io/stargate/sgv2/jsonapi/api/v1/BillingS3ExportIntegrationTest.java +++ b/src/test/java/io/stargate/sgv2/jsonapi/api/v1/BillingS3ExportIntegrationTest.java @@ -40,6 +40,7 @@ * *

Methods are ordered: the last test stops the S3Mock container to prove a failing export never * affects the data API, which kills S3 for the rest of the class — nothing may run after it. + * */ @QuarkusIntegrationTest @WithTestResource(value = DseTestResource.class) From 60b02c8f1f2823ca6dc2febd33f3818b17113b7f Mon Sep 17 00:00:00 2001 From: Hazel Date: Tue, 21 Jul 2026 21:43:03 -0700 Subject: [PATCH 56/65] empty commit --- .../sgv2/jsonapi/api/v1/BillingS3ExportIntegrationTest.java | 1 - 1 file changed, 1 deletion(-) diff --git a/src/test/java/io/stargate/sgv2/jsonapi/api/v1/BillingS3ExportIntegrationTest.java b/src/test/java/io/stargate/sgv2/jsonapi/api/v1/BillingS3ExportIntegrationTest.java index 6e4d835f0b..b121a50540 100644 --- a/src/test/java/io/stargate/sgv2/jsonapi/api/v1/BillingS3ExportIntegrationTest.java +++ b/src/test/java/io/stargate/sgv2/jsonapi/api/v1/BillingS3ExportIntegrationTest.java @@ -40,7 +40,6 @@ * *

Methods are ordered: the last test stops the S3Mock container to prove a failing export never * affects the data API, which kills S3 for the rest of the class — nothing may run after it. - * */ @QuarkusIntegrationTest @WithTestResource(value = DseTestResource.class) From 72de98c6cfc729e4b6994818f890e198628a7350 Mon Sep 17 00:00:00 2001 From: Hazel Date: Tue, 21 Jul 2026 22:46:37 -0700 Subject: [PATCH 57/65] empty commit From 7bccc50744f70e326a97a6c39a10610f99371c23 Mon Sep 17 00:00:00 2001 From: Hazel Date: Tue, 21 Jul 2026 23:05:24 -0700 Subject: [PATCH 58/65] empty commit From a917c3269b447758e648da50dfac8d350d4df63e Mon Sep 17 00:00:00 2001 From: Aaron Morton Date: Wed, 5 Aug 2026 09:47:02 +1200 Subject: [PATCH 59/65] WIP --- .../jsonapi/api/request/RequestContext.java | 2 +- .../sgv2/jsonapi/config/BillingConfig.java | 2 +- .../jsonapi/config/BillingS3ExportConfig.java | 23 +- .../provider => metrics}/BillingMetrics.java | 27 +- .../{provider => billing}/Billing.java | 4 +- .../{provider => billing}/BillingEvent.java | 2 +- .../BillingEventType.java | 2 +- .../jsonapi/service/billing/BillingQueue.java | 194 ++++++++++++ .../BillingS3HandlerInstaller.java | 22 +- .../service/billing/BillingS3LogHandler.java | 284 ++++++++++++++++++ .../{provider => billing}/DefaultBilling.java | 4 +- .../S3BatchUploader.java | 11 +- .../service/provider/BillingQueue.java | 115 ------- .../service/provider/BillingS3LogHandler.java | 248 --------------- .../stargate/sgv2/jsonapi/TestConstants.java | 2 +- .../BillingEventTest.java | 3 +- .../BillingMetricsTest.java | 4 +- .../BillingQueueTest.java | 21 +- .../BillingS3HandlerInstallerTest.java | 7 +- .../BillingS3LogHandlerTest.java | 4 +- .../{provider => billing}/BillingTest.java | 6 +- .../DefaultBillingTest.java | 7 +- .../S3BatchUploaderTest.java | 3 +- 23 files changed, 573 insertions(+), 424 deletions(-) rename src/main/java/io/stargate/sgv2/jsonapi/{service/provider => metrics}/BillingMetrics.java (78%) rename src/main/java/io/stargate/sgv2/jsonapi/service/{provider => billing}/Billing.java (95%) rename src/main/java/io/stargate/sgv2/jsonapi/service/{provider => billing}/BillingEvent.java (98%) rename src/main/java/io/stargate/sgv2/jsonapi/service/{provider => billing}/BillingEventType.java (98%) create mode 100644 src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingQueue.java rename src/main/java/io/stargate/sgv2/jsonapi/service/{provider => billing}/BillingS3HandlerInstaller.java (82%) create mode 100644 src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3LogHandler.java rename src/main/java/io/stargate/sgv2/jsonapi/service/{provider => billing}/DefaultBilling.java (98%) rename src/main/java/io/stargate/sgv2/jsonapi/service/{provider => billing}/S3BatchUploader.java (96%) delete mode 100644 src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingQueue.java delete mode 100644 src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandler.java rename src/test/java/io/stargate/sgv2/jsonapi/service/{provider => billing}/BillingEventTest.java (97%) rename src/test/java/io/stargate/sgv2/jsonapi/service/{provider => billing}/BillingMetricsTest.java (95%) rename src/test/java/io/stargate/sgv2/jsonapi/service/{provider => billing}/BillingQueueTest.java (89%) rename src/test/java/io/stargate/sgv2/jsonapi/service/{provider => billing}/BillingS3HandlerInstallerTest.java (94%) rename src/test/java/io/stargate/sgv2/jsonapi/service/{provider => billing}/BillingS3LogHandlerTest.java (99%) rename src/test/java/io/stargate/sgv2/jsonapi/service/{provider => billing}/BillingTest.java (94%) rename src/test/java/io/stargate/sgv2/jsonapi/service/{provider => billing}/DefaultBillingTest.java (97%) rename src/test/java/io/stargate/sgv2/jsonapi/service/{provider => billing}/S3BatchUploaderTest.java (99%) diff --git a/src/main/java/io/stargate/sgv2/jsonapi/api/request/RequestContext.java b/src/main/java/io/stargate/sgv2/jsonapi/api/request/RequestContext.java index 10ddb2a9d9..29868fe35a 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/api/request/RequestContext.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/api/request/RequestContext.java @@ -15,7 +15,7 @@ import io.stargate.sgv2.jsonapi.config.feature.ApiFeatures; import io.stargate.sgv2.jsonapi.config.feature.FeaturesConfig; import io.stargate.sgv2.jsonapi.logging.LoggingMDCContext; -import io.stargate.sgv2.jsonapi.service.provider.Billing; +import io.stargate.sgv2.jsonapi.service.billing.Billing; import io.stargate.sgv2.jsonapi.service.schema.SchemaRegistry; import io.vertx.ext.web.RoutingContext; import jakarta.enterprise.context.RequestScoped; diff --git a/src/main/java/io/stargate/sgv2/jsonapi/config/BillingConfig.java b/src/main/java/io/stargate/sgv2/jsonapi/config/BillingConfig.java index 08b9985400..58ac7c1e39 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/config/BillingConfig.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/config/BillingConfig.java @@ -2,7 +2,7 @@ import io.smallrye.config.ConfigMapping; import io.smallrye.config.WithDefault; -import io.stargate.sgv2.jsonapi.service.provider.BillingEventType; +import io.stargate.sgv2.jsonapi.service.billing.BillingEventType; import java.util.List; import java.util.Optional; import java.util.Set; diff --git a/src/main/java/io/stargate/sgv2/jsonapi/config/BillingS3ExportConfig.java b/src/main/java/io/stargate/sgv2/jsonapi/config/BillingS3ExportConfig.java index 28dd29e023..4e89d42c4a 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/config/BillingS3ExportConfig.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/config/BillingS3ExportConfig.java @@ -9,7 +9,7 @@ @ConfigMapping(prefix = "stargate.jsonapi.billing.s3") public interface BillingS3ExportConfig { - /** Master switch: when false the export handler is never installed. */ + /** when false the export handler is never installed. */ @WithDefault("false") boolean enabled(); @@ -19,18 +19,23 @@ public interface BillingS3ExportConfig { /** S3 bucket region */ Optional bucketRegion(); - /** Only for non-AWS S3 endpoints (e.g. S3Mock in tests). */ + /** Only for non-AWS S3 endpoints (e.g. S3Mock in tests). + * TODO: XXX EXPLAIN WHAT THIS SHOULD SET SET TO + * */ Optional endpointOverride(); - /** Line-count seal: a buffered batch is shipped once it holds this many events. */ + /** */ @WithDefault("50") - int maxEvents(); + int maxEventsPerBatch(); - /** UTF-8 NDJSON byte seal; a batch may exceed it by one whole event. */ + /** + * Max bytes to include in a batch, NOTE: if a single event is bigger than this it will be sent in a batch still. + * 2097152 == 2 MB + * */ @WithDefault("2097152") - long maxBytes(); + long maxBytesPerBatch(); - /** Age flush period: buffered events are shipped at least this often, sealed or not. */ + /** Age flush period: buffered events are shipped at least this often */ @WithDefault("PT30S") Duration maxAge(); @@ -42,7 +47,9 @@ public interface BillingS3ExportConfig { @WithDefault("4") int uploadConcurrency(); - /** Budget for draining the buffer at shutdown; keep below the pod termination grace period. */ + /** Budget for draining the buffer at shutdown; keep below the pod termination grace period. + * TODO: XXX WHAT IS THE CURRENT TERMINATION PERIOD ? + * */ @WithDefault("PT20S") Duration shutdownTimeout(); } diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingMetrics.java b/src/main/java/io/stargate/sgv2/jsonapi/metrics/BillingMetrics.java similarity index 78% rename from src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingMetrics.java rename to src/main/java/io/stargate/sgv2/jsonapi/metrics/BillingMetrics.java index e2dc493811..a265c809be 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingMetrics.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/metrics/BillingMetrics.java @@ -1,4 +1,4 @@ -package io.stargate.sgv2.jsonapi.service.provider; +package io.stargate.sgv2.jsonapi.metrics; import io.micrometer.core.instrument.Counter; import io.micrometer.core.instrument.Gauge; @@ -25,7 +25,9 @@ public final class BillingMetrics { private static final Logger LOG = LoggerFactory.getLogger(BillingMetrics.class); private static final long DROP_WARN_INTERVAL_NANOS = TimeUnit.MINUTES.toNanos(10); + // Count of the events sent to the S3 logger private final Counter offered; + private final Counter droppedCapacity; private final Counter droppedShutdown; private final Counter flushed; @@ -33,24 +35,25 @@ public final class BillingMetrics { private final Counter batchesUploaded; private final Counter batchesFailed; private final AtomicLong lastDeliveryEpochSeconds = new AtomicLong(0); - private final AtomicLong lastDropWarnNanos; - private final long queueCapacity; + /** * @param depthSource live queue depth, exposed read-only as {@code billing.s3.queue.depth} - * @param queueCapacity quoted in the buffer-full warning */ public BillingMetrics( - MeterRegistry meterRegistry, Supplier depthSource, long queueCapacity) { - this.queueCapacity = queueCapacity; - this.lastDropWarnNanos = new AtomicLong(System.nanoTime() - DROP_WARN_INTERVAL_NANOS); + MeterRegistry meterRegistry, Supplier depthSource) { + + this.offered = meterRegistry.counter("billing.s3.events.offered"); this.droppedCapacity = meterRegistry.counter("billing.s3.events.dropped", "reason", "capacity"); + this.droppedShutdown = meterRegistry.counter("billing.s3.events.dropped", "reason", "shutdown"); + this.flushed = meterRegistry.counter("billing.s3.events.flushed"); this.failed = meterRegistry.counter("billing.s3.events.failed"); this.batchesUploaded = meterRegistry.counter("billing.s3.batches.uploaded"); this.batchesFailed = meterRegistry.counter("billing.s3.batches.failed"); + // Catches stalls with no failures to count (e.g. the flush trigger died): alert on staleness // gated by offered/depth, so idle time isn't mistaken for a dead export. Gauge.builder( @@ -72,16 +75,6 @@ public void recordOffered() { /** A line was dropped on a full buffer; warns rate-limited so a sustained stall stays visible. */ public void recordDropped() { droppedCapacity.increment(); - long now = System.nanoTime(); - long prev = lastDropWarnNanos.get(); - // Rate limit for the buffer-full warning: keeps a drop storm from spamming the logs while - // still surfacing a second stall long after the first. - if (now - prev >= DROP_WARN_INTERVAL_NANOS && lastDropWarnNanos.compareAndSet(prev, now)) { - LOG.warn( - "Billing S3 export backlog full ({} events): shedding billing events because S3 uploads" - + " are slower than ingest. Every shed line is counted by billing.s3.events.dropped.", - queueCapacity); - } } /** Events still buffered when the shutdown budget ran out; close() logs the tombstone. */ diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/Billing.java b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/Billing.java similarity index 95% rename from src/main/java/io/stargate/sgv2/jsonapi/service/provider/Billing.java rename to src/main/java/io/stargate/sgv2/jsonapi/service/billing/Billing.java index 3bb0a72549..3b2a92e2a6 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/Billing.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/Billing.java @@ -1,8 +1,10 @@ -package io.stargate.sgv2.jsonapi.service.provider; +package io.stargate.sgv2.jsonapi.service.billing; import io.stargate.sgv2.jsonapi.config.BillingConfig; import io.stargate.sgv2.jsonapi.config.feature.ApiFeature; import io.stargate.sgv2.jsonapi.config.feature.ApiFeatures; +import io.stargate.sgv2.jsonapi.service.provider.ModelUsage; + import java.util.Objects; /** diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingEvent.java b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingEvent.java similarity index 98% rename from src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingEvent.java rename to src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingEvent.java index 2f2bbcdfa3..19dbe59263 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingEvent.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingEvent.java @@ -1,4 +1,4 @@ -package io.stargate.sgv2.jsonapi.service.provider; +package io.stargate.sgv2.jsonapi.service.billing; import com.fasterxml.jackson.annotation.JsonInclude; import com.fasterxml.jackson.annotation.JsonProperty; diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingEventType.java b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingEventType.java similarity index 98% rename from src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingEventType.java rename to src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingEventType.java index 43f90cb2d0..c888bc1d36 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingEventType.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingEventType.java @@ -1,4 +1,4 @@ -package io.stargate.sgv2.jsonapi.service.provider; +package io.stargate.sgv2.jsonapi.service.billing; import com.fasterxml.jackson.annotation.JsonValue; import java.util.EnumSet; diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingQueue.java b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingQueue.java new file mode 100644 index 0000000000..a34b118712 --- /dev/null +++ b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingQueue.java @@ -0,0 +1,194 @@ +package io.stargate.sgv2.jsonapi.service.billing; + +import com.google.common.annotations.VisibleForTesting; + +import java.time.Duration; +import java.time.Instant; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.concurrent.ArrayBlockingQueue; +import java.util.concurrent.BlockingQueue; +import java.util.concurrent.atomic.AtomicLong; + +/** + * Bounded in-memory buffer that owns the batching policy of the billing S3 export: it decides when + * a batch is sealed ({@code maxEvents} lines or {@code maxBytes} UTF-8 NDJSON bytes) and hands out + * drained {@link Batch}es. A batch may exceed {@code maxBytes} by one whole line; lines are never + * split. + */ +public class BillingQueue { + + private final BlockingQueue queue; + + // Approximate buffered bytes (see lineBytes); + private final AtomicLong queuedBytes = new AtomicLong(0); + private final int maxBatchSize; + private final long maxBytes; + private final Duration maxAge; + + public BillingQueue(int maxBatchSize, long maxBatchBytes, Duration maxAge, int queueCapacity) { + + if (maxBatchSize < 1) { + throw new IllegalArgumentException("maxBatchSize must be >= 1, got: " + maxBatchSize ); + } + if (maxBatchBytes < 1) { + throw new IllegalArgumentException("maxBatchBytes must be >= 1, got: " + maxBatchBytes ); + } + if (queueCapacity < 1) { + throw new IllegalArgumentException("queueCapacity must be >= 1, got: " + queueCapacity); + } + if (maxAge == null || maxAge.isNegative() || maxAge.isZero()) { + throw new IllegalArgumentException("maxAge must be positive, got: " + maxAge); + } + this.maxBatchSize = maxBatchSize; + this.maxBytes = maxBatchBytes; + this.maxAge = maxAge; + + // must be concurrent to handle multiple threads + this.queue = new ArrayBlockingQueue<>(queueCapacity); + } + + public boolean isEmpty() { + return queue.isEmpty(); + } + + public int size() { + return queue.size(); + } + + /** Buffered-bytes counter, exposed for accounting assertions only (no production caller). */ + @VisibleForTesting + long queuedBytes() { + return queuedBytes.get(); + } + + /** + * Appends a line to the billing queue + * + * @param eventAt when the event was logged + * @param line String for the JSON billing blob + * @return true if the billing line was added to the quee for publishing, false otherwise which + * means the queue has reached capacity and we can no longer buffer events + * + * // AI BELOW + * Buffers one line, or returns {@code false} when the capacity bound is hit. + * + * @param eventAt when the event was logged; carried through to {@link Batch#oldestEventAt()} + */ + public boolean offer(Instant eventAt, String line) { + + var newEntry = new Entry(eventAt, line); + if (!queue.offer(newEntry)) { + return false; + } + queuedBytes.addAndGet(newEntry.lineBytes()); + return true; + } + + private Duration oldestEntry() { + var head = queue.peek(); + + return head == null ? Duration.ZERO : Duration.between(head.eventAt(), Instant.now()); + } + + /** + * Tests if we should start a new batch + * */ + private BillingBatchReason startNextBatch() { + + if (queue.size() >= maxBatchSize) { + return BillingBatchReason.MAX_BATCH_SIZE_EXCEEDED + } + if (queuedBytes.get() > maxBytes) { + return BillingBatchReason.MAX_BYTES_EXCEEDED; + } + if (oldestEntry().compareTo(maxAge) >= 0){ + return BillingBatchReason.MAX_AGE_EXCEEDED; + } + return null; + } + + /** + * + * Assumes we are not running async + * OLD BELOW: + * + * Removes and returns up to one sealed batch + * */ + public Batch maybeDrain() { + + var batchReason = startNextBatch(); + if (batchReason == null) { + return null; + } + + List lines = new ArrayList<>(maxBatchSize); + Instant oldestEventAt = null; + long batchBytes = 0; + Entry entry; + + // No matter why we started we create a full batch, e.g. we could start because the oldest + // entry is past maxAge, but we still fill the batch. + while (lines.size() < maxBatchSize && batchBytes < maxBytes && (entry = queue.poll()) != null) { + + if (oldestEventAt == null || entry.eventAt().isBefore(oldestEventAt)) { + oldestEventAt = entry.eventAt(); + } + + lines.add(entry.line()); + var lineBytes = entry.lineBytes(); + queuedBytes.addAndGet(-lineBytes); + batchBytes += lineBytes; + } + + // sanity check, in case of concurrent calls + if (lines.isEmpty()){ + return null; + } + + return new Batch(batchReason, lines, oldestEventAt); + } + + + public enum BillingBatchReason { + MAX_BATCH_SIZE_EXCEEDED, + MAX_BYTES_EXCEEDED, + MAX_AGE_EXCEEDED + } + + /** + * OLD BELOW + * + * One drained, sealed batch. {@code oldestEventAt} is the minimum event time across {@code lines} + * — queue order is enqueue order, not event-time order, under concurrent publish. + */ + public record Batch(BillingBatchReason batchReason, List lines, Instant oldestEventAt) { + + public Batch{ + lines = Collections.unmodifiableList(lines); + } + + public int size(){ + return lines.size(); + } + } + + /** + * Holder for the billing event lines we get called with. + * @param eventAt When the event happened + * @param line The billing event line to record + */ + private record Entry(Instant eventAt, String line) { + + /** + * Gets the length of the line in bytes, + *

+ * Kind of a hack, we are counting unicode code points and calling that 1 byte. Should work + * for ascii text, will undercount if there is non ascii chars but everthing in billing should be ascii + */ + public int lineBytes() { + return line.length() + 1; // +1 is for a newline + } + } +} diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3HandlerInstaller.java b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3HandlerInstaller.java similarity index 82% rename from src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3HandlerInstaller.java rename to src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3HandlerInstaller.java index da0086a72d..5c57548f17 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3HandlerInstaller.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3HandlerInstaller.java @@ -1,4 +1,4 @@ -package io.stargate.sgv2.jsonapi.service.provider; +package io.stargate.sgv2.jsonapi.service.billing; import io.micrometer.core.instrument.MeterRegistry; import io.quarkus.runtime.ShutdownEvent; @@ -11,6 +11,9 @@ import org.slf4j.LoggerFactory; /** + * + * TODO: XXX MAKE THIS COMMENTS READABLE BY A HUMAN + * * Attaches a {@link BillingS3LogHandler} to the {@code billing.events} JUL logger at startup (when * {@link BillingS3ExportConfig#enabled()} is {@code true}) and removes + closes it on shutdown for * a graceful drain. @@ -26,8 +29,7 @@ @ApplicationScoped public class BillingS3HandlerInstaller { - private static final org.slf4j.Logger LOG = - LoggerFactory.getLogger(BillingS3HandlerInstaller.class); + private static final org.slf4j.Logger LOGGER = LoggerFactory.getLogger(BillingS3HandlerInstaller.class); static final String BILLING_LOGGER_NAME = "billing.events"; @@ -43,8 +45,9 @@ public BillingS3HandlerInstaller(BillingS3ExportConfig config, MeterRegistry met } void onStart(@Observes StartupEvent event) { + if (!config.enabled()) { - LOG.debug("Billing S3 export disabled (stargate.jsonapi.billing.s3.enabled=false)"); + LOGGER.debug("Billing S3 export disabled (stargate.jsonapi.billing.s3.enabled=false)"); return; } @@ -55,10 +58,11 @@ void onStart(@Observes StartupEvent event) { var uploader = S3BatchUploader.create(region, bucket, config.endpointOverride()); this.handler = new BillingS3LogHandler(config, uploader, meterRegistry); + // TODO: LOGGER NAME SHOULD BE IN CONFIG Logger.getLogger(BILLING_LOGGER_NAME).addHandler(this.handler); - LOG.info( - "Installed billing S3 export handler on '{}' → bucket '{}' (region '{}', endpointOverride={})", + LOGGER.info( + "Attached billing S3 export handler to logger named {}, bucket={}, region={}, endpointOverride={}", BILLING_LOGGER_NAME, bucket, region, @@ -66,16 +70,20 @@ void onStart(@Observes StartupEvent event) { } void onStop(@Observes ShutdownEvent event) { + if (this.handler == null) { return; } + + // TODO: XXX WHY DO THIS ? Logger.getLogger(BILLING_LOGGER_NAME).removeHandler(this.handler); + // close() isn't expected to throw, but if it does (e.g. client.close() failing), letting it // propagate would disrupt other components' cleanup in Quarkus's shutdown sequence. try { this.handler.close(); } catch (Exception e) { - LOG.warn("Error during billing S3 export handler shutdown", e); + LOGGER.warn("Error during billing S3 export handler shutdown", e); } } } diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3LogHandler.java b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3LogHandler.java new file mode 100644 index 0000000000..88cc82f123 --- /dev/null +++ b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3LogHandler.java @@ -0,0 +1,284 @@ +package io.stargate.sgv2.jsonapi.service.billing; + +import com.google.common.annotations.VisibleForTesting; +import io.smallrye.mutiny.Uni; + +import java.time.Duration; +import java.util.Objects; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.logging.Handler; +import java.util.logging.LogRecord; + +import io.stargate.sgv2.jsonapi.metrics.BillingMetrics; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * A Logging handler designed to be used wioth the Billing system. It accpets billing event log + * messges, batches them, and then sends to S3. + *

+ * See {@link BillingS3HandlerInstaller} for setup. + *

+ * + * // AI SLOP BELOW + * JUL handler that turns {@code billing.events} log lines into batched S3 objects. + * + *

Division of labor: {@link BillingQueue} decides when a batch seals, {@link AsyncBatchUploader} + * decides what an S3 object looks like, and this class decides when uploads run — the flush + * triggers (seal on publish, age tick, drain on close), the upload-concurrency gate, and metrics. + * + *

Delivery is at-most-once by design: publish never waits for queue capacity, full buffers drop + * new lines, and close drains best-effort within {@code shutdownTimeout}. + */ +public final class BillingS3LogHandler extends Handler { + + private static final Logger LOGGER = LoggerFactory.getLogger(BillingS3LogHandler.class); + + private final Object wakeupSignal = new Object(); + private final AtomicBoolean isClosed = new AtomicBoolean(false); + private final AtomicBoolean isDraining = new AtomicBoolean(false); + private final CountDownLatch finishedLatch = new CountDownLatch(1); + + private final AsyncBatchUploader uploader; + private final BillingMetrics billingMetrics; + private final BillingQueue billingQueue; + + @VisibleForTesting + BillingS3LogHandler( + AsyncBatchUploader uploader, + BillingQueue billingQueue, + BillingMetrics billingMetrics) { + + this.billingQueue = billingQueue; + this.uploader = uploader; + this.billingMetrics = Objects.requireNonNull(billingMetrics); + } + + private static Duration requirePositive(Duration value, String property) { + if (value == null || value.isNegative() || value.isZero()) { + throw new IllegalArgumentException( + "stargate.jsonapi.billing.s3." + property + " must be > 0 (was " + value + ")"); + } + return value; + } + + // ============================================================ + // Overrides for java.util.logging.Handler + // ============================================================ + + @Override + public void publish(LogRecord record) { + + if (record == null || isClosed.get()) { + return; + } + + // TODO: XXX : WHAT DOES MEAN "This handler never runs a Formatter, so a parameterized call would ship its" + + // Producer contract (DefaultBilling): the message is the final JSON line, logged without {} + // placeholders. This handler never runs a Formatter, so a parameterized call would ship its + // raw template. getInstant() is when the producer logged it — within microseconds of the + // "timestamp" it embedded in the JSON, and the object key only needs minute resolution. + + var line = record.getMessage(); + if (line == null || line.isBlank()) { + return; + } + + billingMetrics.recordOffered(); + if (!billingQueue.offer(record.getInstant(), line)) { + // Bounded buffer full: drop and count + billingMetrics.recordDropped(); + } + // flushing runs every second so nothing more to do + } + + @Override + public void flush() { + setWakeupSignal(); + } + + /** + * Drains what remains through the normal flush pipeline, bounded by {@code shutdownTimeout}. The + * budget only bites when S3 is already failing: it converts a silent SIGKILL into a logged count + * of abandoned events and lets the rest of shutdown proceed. + */ + @Override + public void close() { + + isClosed.set(true); + isDraining.set(true); + flush(); + + try { + if (!finishedLatch.await(30, TimeUnit.SECONDS)) { + LOGGER.warn("close() - Billing upload loop did not stop within 30s, interrupting"); + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + LOGGER.warn("close() - Interrupted waiting for billing upload loop to finish"); + } + finally { + uploader.close(); + } + + } + + // ============================================================ + // Flush pipeline + // ============================================================ + + private void setWakeupSignal() { + synchronized (wakeupSignal) { + wakeupSignal.notifyAll(); + } + } + + void startPublishing() { + + BillingQueue.Batch batch; + try { + while (true) { + + synchronized (wakeupSignal) { + if (!isDraining.get()) { + try { + wakeupSignal.wait(1000); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + return; + } + } + } + + while ((batch = billingQueue.maybeDrain()) != null) { + // indefinitely() is bounded by the 10s ifNoItem() timeout inside deferBatch(), + // and failures are recovered there, so this neither hangs nor throws. + deferBatch(batch).await().indefinitely(); + } + + if (isDraining.get()){ + return; + } + } + } + finally { + billingMetrics.recordAbandonedAtShutdown(billingQueue.size()); + if (!billingQueue.isEmpty()){ + LOGGER.warn("start() - finished with abandoned billing events, billingQueue.size():{} " , billingQueue.size()); + } + finishedLatch.countDown(); + } + } + + private Uni deferBatch(BillingQueue.Batch batch) { + + // upload() is called at subscription, not when this method returns. + // deferred also converts a synchronous throw from upload() into a Uni failure. + + return Uni.createFrom() + .deferred(() -> uploader.upload(batch)) + .ifNoItem() + .after(Duration.ofSeconds(10)) + .fail() + .onItemOrFailure() + .invoke( + (item, failure) -> { + if (failure != null) { + billingMetrics.recordBatchFailed(batch.size()); + LOGGER.error("Failed to upload billing S3 batch ({} events)", batch.size(), failure); + } else { + billingMetrics.recordBatchDelivered(batch.size()); + } + }) + .onFailure() + .recoverWithNull(); + } + +// /** Seal-triggered flush: ship when the buffer has a full batch by count or bytes. */ +// private void maybeFlush() { +// if (eventQueue.shouldFlush()) { +// tryFlush(); +// } +// } + +// /** +// * Age trigger: every {@code maxAge} tick ships whatever is buffered, sealed or not. Deliberately +// * no head-age check: flushing only entries older than {@code maxAge} would let an event that just +// * missed a tick wait ~2x{@code maxAge}, while shipping unconditionally bounds every wait by one +// * period — at the cost of an occasional small object when a tick lands just after a seal flush. +// * +// *

Catches everything: an escaped throwable would silently cancel all future runs of a +// * fixed-rate task. +// */ +// @VisibleForTesting +// void onAgeTick() { +// try { +// if (!eventQueue.isEmpty()) { +// tryFlush(); +// } +// } catch (Throwable t) { +// LOG.error("Billing S3 export age-flush tick failed", t); +// } +// } +// +// /** +// * Claims an in-flight slot (non-blocking CAS, at most {@link #uploadConcurrency} held) and, on +// * success, drains + uploads one batch asynchronously. When the upload settles the slot is +// * released and the seal condition re-checked: a full batch may have accumulated meanwhile. +// */ +// private void tryFlush() { +// +// int prev = inFlight.getAndUpdate(n -> n < uploadConcurrency ? n + 1 : n); +// if (prev >= uploadConcurrency) { +// return; +// } +// Uni.createFrom() +// .item(eventQueue::maybeDrain) +// .runSubscriptionOn(Infrastructure.getDefaultWorkerPool()) +// .flatMap(this::uploadBatch) +// .eventually( +// () -> { +// inFlight.getAndDecrement(); +// maybeFlush(); +// }) +// .subscribe() +// .with(ignored -> {}, failure -> LOG.error("Billing S3 export flush failed", failure)); +// } + + /** Uploads one batch; never fails the pipeline — a batch that exhausts retries is counted. */ +// private Uni uploadBatch(BillingQueue.Batch batch) { +// if (batch.isEmpty()) { +// return Uni.createFrom().voidItem(); +// } +// int size = batch.size(); +// // Runs immediately on subscription; deferred only turns a throw before upload() returns a Uni +// // into a Uni failure handled below. +// return Uni.createFrom() +// .deferred(() -> uploader.upload(batch)) +// .onItem() +// .invoke(() -> billingMetrics.recordBatchDelivered(size)) +// .onFailure() +// .invoke(t -> LOG.error("Failed to upload billing S3 batch ({} events)", size, t)) +// .onFailure() +// .recoverWithItem( +// () -> { +// billingMetrics.recordBatchFailed(size); +// return null; +// }); +// } + + /** + * Uploads one sealed batch to the export destination; owns the object key and body encoding. + * Implementations must tolerate concurrent calls. + */ + @FunctionalInterface + public interface AsyncBatchUploader extends AutoCloseable { + Uni upload(BillingQueue.Batch batch); + + @Override + default void close() {} + } +} diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/DefaultBilling.java b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/DefaultBilling.java similarity index 98% rename from src/main/java/io/stargate/sgv2/jsonapi/service/provider/DefaultBilling.java rename to src/main/java/io/stargate/sgv2/jsonapi/service/billing/DefaultBilling.java index 9ab810bd4a..f149226ab3 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/DefaultBilling.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/DefaultBilling.java @@ -1,4 +1,4 @@ -package io.stargate.sgv2.jsonapi.service.provider; +package io.stargate.sgv2.jsonapi.service.billing; import static io.stargate.sgv2.jsonapi.util.StringUtil.requireNonBlank; @@ -14,6 +14,8 @@ import java.util.Objects; import java.util.Set; import java.util.UUID; + +import io.stargate.sgv2.jsonapi.service.provider.ModelUsage; import org.slf4j.Logger; import org.slf4j.LoggerFactory; diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/S3BatchUploader.java b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/S3BatchUploader.java similarity index 96% rename from src/main/java/io/stargate/sgv2/jsonapi/service/provider/S3BatchUploader.java rename to src/main/java/io/stargate/sgv2/jsonapi/service/billing/S3BatchUploader.java index dd29b535b2..8cf1758086 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/S3BatchUploader.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/S3BatchUploader.java @@ -1,4 +1,4 @@ -package io.stargate.sgv2.jsonapi.service.provider; +package io.stargate.sgv2.jsonapi.service.billing; import com.google.common.annotations.VisibleForTesting; import io.smallrye.mutiny.Uni; @@ -12,6 +12,7 @@ import java.util.Objects; import java.util.Optional; import java.util.UUID; + import software.amazon.awssdk.core.async.AsyncRequestBody; import software.amazon.awssdk.regions.Region; import software.amazon.awssdk.services.s3.S3AsyncClient; @@ -42,11 +43,15 @@ public class S3BatchUploader implements BillingS3LogHandler.AsyncBatchUploader { } public static S3BatchUploader create( + String region, String bucket, Optional endpointOverride) { - if (region == null || region.isBlank()) + if (region == null || region.isBlank()) { throw new IllegalArgumentException("stargate.jsonapi.billing.s3.bucket-region must be set"); - if (bucket == null || bucket.isBlank()) + } + + if (bucket == null || bucket.isBlank()) { throw new IllegalArgumentException("stargate.jsonapi.billing.s3.bucket must be set"); + } Objects.requireNonNull(endpointOverride, "endpointOverride must not be null"); // Credentials resolve from the SDK's default provider chain (env vars, web-identity/OIDC diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingQueue.java b/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingQueue.java deleted file mode 100644 index 33171df5c9..0000000000 --- a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingQueue.java +++ /dev/null @@ -1,115 +0,0 @@ -package io.stargate.sgv2.jsonapi.service.provider; - -import com.google.common.annotations.VisibleForTesting; -import java.time.Instant; -import java.util.ArrayList; -import java.util.List; -import java.util.concurrent.ArrayBlockingQueue; -import java.util.concurrent.BlockingQueue; -import java.util.concurrent.atomic.AtomicLong; - -/** - * Bounded in-memory buffer that owns the batching policy of the billing S3 export: it decides when - * a batch is sealed ({@code maxEvents} lines or {@code maxBytes} UTF-8 NDJSON bytes) and hands out - * drained {@link Batch}es. A batch may exceed {@code maxBytes} by one whole line; lines are never - * split. - */ -public final class BillingQueue { - - private final BlockingQueue queue; - // Approximate buffered NDJSON bytes (see lineBytes); - private final AtomicLong queuedBytes = new AtomicLong(0); - private final int batchSize; - private final long maxBytes; - - public BillingQueue(int maxEvents, long maxBytes, int queueCapacity) { - if (maxEvents < 1) { - throw new IllegalArgumentException( - "stargate.jsonapi.billing.s3.max-events must be >= 1 (was " + maxEvents + ")"); - } - if (maxBytes < 1) { - throw new IllegalArgumentException( - "stargate.jsonapi.billing.s3.max-bytes must be >= 1 (was " + maxBytes + ")"); - } - if (queueCapacity < 1) { - throw new IllegalArgumentException( - "stargate.jsonapi.billing.s3.queue-capacity must be >= 1 (was " + queueCapacity + ")"); - } - this.batchSize = maxEvents; - this.maxBytes = maxBytes; - this.queue = new ArrayBlockingQueue<>(queueCapacity); - } - - /** - * Buffers one line, or returns {@code false} when the capacity bound is hit. - * - * @param eventAt when the event was logged; carried through to {@link Batch#oldestEventAt()} - */ - public boolean offer(Instant eventAt, String line) { - if (!queue.offer(new Entry(eventAt, line))) { - return false; - } - queuedBytes.addAndGet(lineBytes(line)); - return true; - } - - /** True once a seal is reached: a full batch by line count, or {@code maxBytes} buffered. */ - public boolean shouldFlush() { - return queue.size() >= batchSize || queuedBytes.get() >= maxBytes; - } - - public boolean isEmpty() { - return queue.isEmpty(); - } - - public int size() { - return queue.size(); - } - - /** Buffered-bytes counter, exposed for accounting assertions only (no production caller). */ - @VisibleForTesting - long queuedBytes() { - return queuedBytes.get(); - } - - /** Removes and returns up to one sealed batch (possibly partial, possibly {@code EMPTY}). */ - public Batch drain() { - List lines = new ArrayList<>(batchSize); - Instant oldestEventAt = null; - long bytes = 0; - Entry entry; - while (lines.size() < batchSize && bytes < maxBytes && (entry = queue.poll()) != null) { - if (oldestEventAt == null || entry.eventAt().isBefore(oldestEventAt)) { - oldestEventAt = entry.eventAt(); - } - lines.add(entry.line()); - bytes += lineBytes(entry.line()); - } - queuedBytes.addAndGet(-bytes); - return oldestEventAt == null ? Batch.EMPTY : new Batch(lines, oldestEventAt); - } - - // String.length() (UTF-16 units) + newline as a cheap stand-in for UTF-8 bytes: exact for the - // ASCII JSON , an undercount for non-ASCII - private static int lineBytes(String line) { - return line.length() + 1; - } - - /** - * One drained, sealed batch. {@code oldestEventAt} is the minimum event time across {@code lines} - * — queue order is enqueue order, not event-time order, under concurrent publish. - */ - public record Batch(List lines, Instant oldestEventAt) { - static final Batch EMPTY = new Batch(List.of(), Instant.EPOCH); - - boolean isEmpty() { - return lines.isEmpty(); - } - - int size() { - return lines.size(); - } - } - - private record Entry(Instant eventAt, String line) {} -} diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandler.java b/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandler.java deleted file mode 100644 index 97af685d37..0000000000 --- a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandler.java +++ /dev/null @@ -1,248 +0,0 @@ -package io.stargate.sgv2.jsonapi.service.provider; - -import com.google.common.annotations.VisibleForTesting; -import io.micrometer.core.instrument.MeterRegistry; -import io.smallrye.mutiny.Uni; -import io.smallrye.mutiny.infrastructure.Infrastructure; -import io.stargate.sgv2.jsonapi.config.BillingS3ExportConfig; -import java.time.Duration; -import java.util.concurrent.ScheduledFuture; -import java.util.concurrent.TimeUnit; -import java.util.concurrent.atomic.AtomicInteger; -import java.util.concurrent.locks.LockSupport; -import java.util.logging.Handler; -import java.util.logging.LogRecord; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; - -/** - * JUL handler that turns {@code billing.events} log lines into batched S3 objects. - * - *

Division of labor: {@link BillingQueue} decides when a batch seals, {@link AsyncBatchUploader} - * decides what an S3 object looks like, and this class decides when uploads run — the flush - * triggers (seal on publish, age tick, drain on close), the upload-concurrency gate, and metrics. - * - *

Delivery is at-most-once by design: publish never waits for queue capacity, full buffers drop - * new lines, and close drains best-effort within {@code shutdownTimeout}. - */ -public final class BillingS3LogHandler extends Handler { - - private static final Logger LOG = LoggerFactory.getLogger(BillingS3LogHandler.class); - - // ---- Collaborators ---- - private final AsyncBatchUploader uploader; - private final BillingMetrics metrics; - private final BillingQueue buffer; - - // ---- Flush pipeline ---- - private final int uploadConcurrency; // max flushes (S3 PUTs) in flight at once - private final AtomicInteger inFlight = new AtomicInteger(0); - private final ScheduledFuture ageFlushTask; - - private final Duration shutdownTimeout; - - public BillingS3LogHandler( - BillingS3ExportConfig config, AsyncBatchUploader uploader, MeterRegistry meterRegistry) { - this( - uploader, - meterRegistry, - config.maxEvents(), - config.maxBytes(), - config.maxAge(), - config.queueCapacity(), - config.uploadConcurrency(), - config.shutdownTimeout()); - } - - @VisibleForTesting - BillingS3LogHandler( - AsyncBatchUploader uploader, - MeterRegistry meterRegistry, - int maxEvents, - long maxBytes, - Duration maxAge, - int queueCapacity, - int uploadConcurrency, - Duration shutdownTimeout) { - if (uploadConcurrency < 1) { - throw new IllegalArgumentException( - "stargate.jsonapi.billing.s3.upload-concurrency must be >= 1 (was " - + uploadConcurrency - + ")"); - } - requirePositive(maxAge, "max-age"); - requirePositive(shutdownTimeout, "shutdown-timeout"); - this.uploader = uploader; - this.uploadConcurrency = uploadConcurrency; - this.shutdownTimeout = shutdownTimeout; - this.buffer = new BillingQueue(maxEvents, maxBytes, queueCapacity); - this.metrics = new BillingMetrics(meterRegistry, buffer::size, queueCapacity); - this.ageFlushTask = - Infrastructure.getDefaultWorkerPool() - .scheduleAtFixedRate( - this::onAgeTick, maxAge.toMillis(), maxAge.toMillis(), TimeUnit.MILLISECONDS); - } - - private static void requirePositive(Duration value, String property) { - if (value == null || value.isNegative() || value.isZero()) { - throw new IllegalArgumentException( - "stargate.jsonapi.billing.s3." + property + " must be > 0 (was " + value + ")"); - } - } - - // ============================================================ - // java.util.logging.Handler - // ============================================================ - - @Override - public void publish(LogRecord record) { - if (record == null) { - return; - } - // Producer contract (DefaultBilling): the message is the final JSON line, logged without {} - // placeholders. This handler never runs a Formatter, so a parameterized call would ship its - // raw template. getInstant() is when the producer logged it — within microseconds of the - // "timestamp" it embedded in the JSON, and the object key only needs minute resolution. - String line = record.getMessage(); - if (line == null || line.isBlank()) { - return; - } - metrics.recordOffered(); - if (!buffer.offer(record.getInstant(), line)) { - // Bounded buffer full: drop and count - metrics.recordDropped(); - return; - } - maybeFlush(); - } - - @Override - public void flush() { - // No-op: shipping is seal-triggered (publish) and age-triggered (tick); close() drains. - } - - /** - * Drains what remains through the normal flush pipeline, bounded by {@code shutdownTimeout}. The - * budget only bites when S3 is already failing: it converts a silent SIGKILL into a logged count - * of abandoned events and lets the rest of shutdown proceed. - */ - @Override - public void close() { - // Don't interrupt a tick already running; the in-flight wait below covers it. - ageFlushTask.cancel(false); - long deadlineNanos = System.nanoTime() + shutdownTimeout.toNanos(); - - // Pump the pipeline until the buffer is drained: tryFlush() is a no-op while all slots are - // busy, and every settled upload frees a slot for the next batch. - while (!buffer.isEmpty() && System.nanoTime() < deadlineNanos) { - tryFlush(); - LockSupport.parkNanos(TimeUnit.MILLISECONDS.toNanos(1)); - } - // Let in-flight uploads (ours and any started before close) settle within the budget. - while (inFlight.get() > 0 && System.nanoTime() < deadlineNanos) { - LockSupport.parkNanos(TimeUnit.MILLISECONDS.toNanos(10)); - } - - int queuedAbandoned = buffer.size(); - int inFlightAbandoned = inFlight.get(); - if (queuedAbandoned > 0 || inFlightAbandoned > 0) { - // Only queued events are counted: in-flight ones settle as failed when the abort lands. - metrics.recordAbandonedAtShutdown(queuedAbandoned); - LOG.warn( - "Billing S3 export shutdown budget ({}) exhausted: dropping {} buffered events," - + " abandoning {} in-flight uploads", - shutdownTimeout, - queuedAbandoned, - inFlightAbandoned); - } - uploader.close(); // aborts anything still in flight - } - - // ============================================================ - // Flush pipeline - // ============================================================ - - /** Seal-triggered flush: ship when the buffer has a full batch by count or bytes. */ - private void maybeFlush() { - if (buffer.shouldFlush()) { - tryFlush(); - } - } - - /** - * Age trigger: every {@code maxAge} tick ships whatever is buffered, sealed or not. Deliberately - * no head-age check: flushing only entries older than {@code maxAge} would let an event that just - * missed a tick wait ~2x{@code maxAge}, while shipping unconditionally bounds every wait by one - * period — at the cost of an occasional small object when a tick lands just after a seal flush. - * - *

Catches everything: an escaped throwable would silently cancel all future runs of a - * fixed-rate task. - */ - @VisibleForTesting - void onAgeTick() { - try { - if (!buffer.isEmpty()) { - tryFlush(); - } - } catch (Throwable t) { - LOG.error("Billing S3 export age-flush tick failed", t); - } - } - - /** - * Claims an in-flight slot (non-blocking CAS, at most {@link #uploadConcurrency} held) and, on - * success, drains + uploads one batch asynchronously. When the upload settles the slot is - * released and the seal condition re-checked: a full batch may have accumulated meanwhile. - */ - private void tryFlush() { - int prev = inFlight.getAndUpdate(n -> n < uploadConcurrency ? n + 1 : n); - if (prev >= uploadConcurrency) { - return; - } - Uni.createFrom() - .item(buffer::drain) - .runSubscriptionOn(Infrastructure.getDefaultWorkerPool()) - .flatMap(this::uploadBatch) - .eventually( - () -> { - inFlight.getAndDecrement(); - maybeFlush(); - }) - .subscribe() - .with(ignored -> {}, failure -> LOG.error("Billing S3 export flush failed", failure)); - } - - /** Uploads one batch; never fails the pipeline — a batch that exhausts retries is counted. */ - private Uni uploadBatch(BillingQueue.Batch batch) { - if (batch.isEmpty()) { - return Uni.createFrom().voidItem(); - } - int size = batch.size(); - // Runs immediately on subscription; deferred only turns a throw before upload() returns a Uni - // into a Uni failure handled below. - return Uni.createFrom() - .deferred(() -> uploader.upload(batch)) - .onItem() - .invoke(() -> metrics.recordBatchDelivered(size)) - .onFailure() - .invoke(t -> LOG.error("Failed to upload billing S3 batch ({} events)", size, t)) - .onFailure() - .recoverWithItem( - () -> { - metrics.recordBatchFailed(size); - return null; - }); - } - - /** - * Uploads one sealed batch to the export destination; owns the object key and body encoding. - * Implementations must tolerate concurrent calls. - */ - @FunctionalInterface - public interface AsyncBatchUploader extends AutoCloseable { - Uni upload(BillingQueue.Batch batch); - - @Override - default void close() {} - } -} diff --git a/src/test/java/io/stargate/sgv2/jsonapi/TestConstants.java b/src/test/java/io/stargate/sgv2/jsonapi/TestConstants.java index e0c42122bc..35be0bddc5 100644 --- a/src/test/java/io/stargate/sgv2/jsonapi/TestConstants.java +++ b/src/test/java/io/stargate/sgv2/jsonapi/TestConstants.java @@ -22,7 +22,7 @@ import io.stargate.sgv2.jsonapi.service.cqldriver.executor.*; import io.stargate.sgv2.jsonapi.service.embedding.operation.EmbeddingProvider; import io.stargate.sgv2.jsonapi.service.embedding.operation.EmbeddingProviderFactory; -import io.stargate.sgv2.jsonapi.service.provider.Billing; +import io.stargate.sgv2.jsonapi.service.billing.Billing; import io.stargate.sgv2.jsonapi.service.reranking.operation.RerankingProviderFactory; import io.stargate.sgv2.jsonapi.service.schema.*; import io.stargate.sgv2.jsonapi.service.schema.collections.CollectionLexicalDefSchemaFactory; diff --git a/src/test/java/io/stargate/sgv2/jsonapi/service/provider/BillingEventTest.java b/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingEventTest.java similarity index 97% rename from src/test/java/io/stargate/sgv2/jsonapi/service/provider/BillingEventTest.java rename to src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingEventTest.java index dbeca3f937..f01a95ea97 100644 --- a/src/test/java/io/stargate/sgv2/jsonapi/service/provider/BillingEventTest.java +++ b/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingEventTest.java @@ -1,10 +1,11 @@ -package io.stargate.sgv2.jsonapi.service.provider; +package io.stargate.sgv2.jsonapi.service.billing; import static net.javacrumbs.jsonunit.JsonAssert.assertJsonEquals; import com.fasterxml.jackson.databind.ObjectMapper; import java.time.Instant; import java.util.UUID; + import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.EnumSource; diff --git a/src/test/java/io/stargate/sgv2/jsonapi/service/provider/BillingMetricsTest.java b/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingMetricsTest.java similarity index 95% rename from src/test/java/io/stargate/sgv2/jsonapi/service/provider/BillingMetricsTest.java rename to src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingMetricsTest.java index ade91aac03..47de9e45c6 100644 --- a/src/test/java/io/stargate/sgv2/jsonapi/service/provider/BillingMetricsTest.java +++ b/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingMetricsTest.java @@ -1,10 +1,12 @@ -package io.stargate.sgv2.jsonapi.service.provider; +package io.stargate.sgv2.jsonapi.service.billing; import static org.assertj.core.api.Assertions.assertThat; import io.micrometer.core.instrument.simple.SimpleMeterRegistry; import java.time.Instant; import java.util.concurrent.atomic.AtomicInteger; + +import io.stargate.sgv2.jsonapi.metrics.BillingMetrics; import org.junit.jupiter.api.Test; /** Guards the meter names and tags — dashboards and alerts key on these exact series. */ diff --git a/src/test/java/io/stargate/sgv2/jsonapi/service/provider/BillingQueueTest.java b/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingQueueTest.java similarity index 89% rename from src/test/java/io/stargate/sgv2/jsonapi/service/provider/BillingQueueTest.java rename to src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingQueueTest.java index 5d0cd55f5a..98cea178dd 100644 --- a/src/test/java/io/stargate/sgv2/jsonapi/service/provider/BillingQueueTest.java +++ b/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingQueueTest.java @@ -1,4 +1,4 @@ -package io.stargate.sgv2.jsonapi.service.provider; +package io.stargate.sgv2.jsonapi.service.billing; import static org.assertj.core.api.Assertions.assertThat; @@ -14,6 +14,7 @@ import java.util.concurrent.Future; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; + import org.junit.jupiter.api.Test; /** Unit tests for {@link BillingQueue}: seal thresholds, drain limits, and batch metadata. */ @@ -49,9 +50,9 @@ void drainStopsAtMaxEventsAndLeavesTheRemainder() { queue.offer(T0, "b"); queue.offer(T0, "c"); - assertThat(queue.drain().lines()).containsExactly("a", "b"); - assertThat(queue.drain().lines()).containsExactly("c"); - assertThat(queue.drain().isEmpty()).isTrue(); + assertThat(queue.maybeDrain().lines()).containsExactly("a", "b"); + assertThat(queue.maybeDrain().lines()).containsExactly("c"); + assertThat(queue.maybeDrain().isEmpty()).isTrue(); } @Test @@ -61,9 +62,9 @@ void drainStopsAtMaxBytesAndLeavesTheRemainder() { queue.offer(T0, "bbbb"); queue.offer(T0, "cccc"); - assertThat(queue.drain().lines()).containsExactly("aaaa", "bbbb"); - assertThat(queue.drain().lines()).containsExactly("cccc"); - assertThat(queue.drain().isEmpty()).isTrue(); + assertThat(queue.maybeDrain().lines()).containsExactly("aaaa", "bbbb"); + assertThat(queue.maybeDrain().lines()).containsExactly("cccc"); + assertThat(queue.maybeDrain().isEmpty()).isTrue(); } @Test @@ -73,7 +74,7 @@ void oldestEventAtIsTheMinimumAcrossTheBatchNotTheHead() { queue.offer(T0.plusSeconds(5), "enqueued-first-but-newer"); queue.offer(T0, "enqueued-second-but-older"); - assertThat(queue.drain().oldestEventAt()).isEqualTo(T0); + assertThat(queue.maybeDrain().oldestEventAt()).isEqualTo(T0); } @Test @@ -99,7 +100,7 @@ void concurrentOfferAndDrainKeepsAccountingConsistent() throws Exception { new Thread( () -> { while (!producersDone.get() || !queue.isEmpty()) { - var batch = queue.drain(); + var batch = queue.maybeDrain(); if (batch.isEmpty()) { Thread.onSpinWait(); } else { @@ -154,7 +155,7 @@ void byteSealResetsOnceDrained() { queue.offer(T0, "bbbb"); assertThat(queue.shouldFlush()).isTrue(); - queue.drain(); + queue.maybeDrain(); assertThat(queue.isEmpty()).isTrue(); assertThat(queue.shouldFlush()).isFalse(); // queuedBytes went back down with the drain diff --git a/src/test/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3HandlerInstallerTest.java b/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3HandlerInstallerTest.java similarity index 94% rename from src/test/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3HandlerInstallerTest.java rename to src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3HandlerInstallerTest.java index f085830f6f..33f85ebf1b 100644 --- a/src/test/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3HandlerInstallerTest.java +++ b/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3HandlerInstallerTest.java @@ -1,4 +1,4 @@ -package io.stargate.sgv2.jsonapi.service.provider; +package io.stargate.sgv2.jsonapi.service.billing; import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThatThrownBy; @@ -13,6 +13,7 @@ import java.util.Arrays; import java.util.Optional; import java.util.logging.Logger; + import org.junit.jupiter.api.Test; /** @@ -28,8 +29,8 @@ private static BillingS3ExportConfig config(boolean enabled, String bucket, Stri when(config.bucket()).thenReturn(Optional.ofNullable(bucket)); when(config.bucketRegion()).thenReturn(Optional.ofNullable(region)); when(config.endpointOverride()).thenReturn(Optional.empty()); - when(config.maxEvents()).thenReturn(50); - when(config.maxBytes()).thenReturn(2_097_152L); + when(config.maxEventsPerBatch()).thenReturn(50); + when(config.maxBytesPerBatch()).thenReturn(2_097_152L); when(config.maxAge()).thenReturn(Duration.ofSeconds(30)); when(config.queueCapacity()).thenReturn(100); when(config.uploadConcurrency()).thenReturn(2); diff --git a/src/test/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandlerTest.java b/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3LogHandlerTest.java similarity index 99% rename from src/test/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandlerTest.java rename to src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3LogHandlerTest.java index e4b085094f..6589edad32 100644 --- a/src/test/java/io/stargate/sgv2/jsonapi/service/provider/BillingS3LogHandlerTest.java +++ b/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3LogHandlerTest.java @@ -1,4 +1,4 @@ -package io.stargate.sgv2.jsonapi.service.provider; +package io.stargate.sgv2.jsonapi.service.billing; import static org.assertj.core.api.Assertions.assertThat; import static org.awaitility.Awaitility.await; @@ -26,6 +26,8 @@ import java.util.concurrent.atomic.AtomicInteger; import java.util.logging.Level; import java.util.logging.LogRecord; + +import io.stargate.sgv2.jsonapi.metrics.BillingMetrics; import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.Test; import org.slf4j.Logger; diff --git a/src/test/java/io/stargate/sgv2/jsonapi/service/provider/BillingTest.java b/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingTest.java similarity index 94% rename from src/test/java/io/stargate/sgv2/jsonapi/service/provider/BillingTest.java rename to src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingTest.java index 399155d241..e5b589cae9 100644 --- a/src/test/java/io/stargate/sgv2/jsonapi/service/provider/BillingTest.java +++ b/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingTest.java @@ -1,4 +1,4 @@ -package io.stargate.sgv2.jsonapi.service.provider; +package io.stargate.sgv2.jsonapi.service.billing; import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThatCode; @@ -12,6 +12,10 @@ import io.stargate.sgv2.jsonapi.config.feature.ApiFeature; import io.stargate.sgv2.jsonapi.config.feature.ApiFeatures; import io.stargate.sgv2.jsonapi.config.feature.FeaturesConfig; +import io.stargate.sgv2.jsonapi.service.provider.ModelInputType; +import io.stargate.sgv2.jsonapi.service.provider.ModelProvider; +import io.stargate.sgv2.jsonapi.service.provider.ModelType; +import io.stargate.sgv2.jsonapi.service.provider.ModelUsage; import io.vertx.core.MultiMap; import java.util.List; import java.util.Map; diff --git a/src/test/java/io/stargate/sgv2/jsonapi/service/provider/DefaultBillingTest.java b/src/test/java/io/stargate/sgv2/jsonapi/service/billing/DefaultBillingTest.java similarity index 97% rename from src/test/java/io/stargate/sgv2/jsonapi/service/provider/DefaultBillingTest.java rename to src/test/java/io/stargate/sgv2/jsonapi/service/billing/DefaultBillingTest.java index 7657c48282..ac207bafe7 100644 --- a/src/test/java/io/stargate/sgv2/jsonapi/service/provider/DefaultBillingTest.java +++ b/src/test/java/io/stargate/sgv2/jsonapi/service/billing/DefaultBillingTest.java @@ -1,4 +1,4 @@ -package io.stargate.sgv2.jsonapi.service.provider; +package io.stargate.sgv2.jsonapi.service.billing; import static java.util.logging.Logger.getLogger; import static net.javacrumbs.jsonunit.JsonAssert.assertJsonEquals; @@ -20,6 +20,11 @@ import java.util.logging.Handler; import java.util.logging.LogRecord; import java.util.stream.Stream; + +import io.stargate.sgv2.jsonapi.service.provider.ModelInputType; +import io.stargate.sgv2.jsonapi.service.provider.ModelProvider; +import io.stargate.sgv2.jsonapi.service.provider.ModelType; +import io.stargate.sgv2.jsonapi.service.provider.ModelUsage; import org.junit.jupiter.api.Test; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.Arguments; diff --git a/src/test/java/io/stargate/sgv2/jsonapi/service/provider/S3BatchUploaderTest.java b/src/test/java/io/stargate/sgv2/jsonapi/service/billing/S3BatchUploaderTest.java similarity index 99% rename from src/test/java/io/stargate/sgv2/jsonapi/service/provider/S3BatchUploaderTest.java rename to src/test/java/io/stargate/sgv2/jsonapi/service/billing/S3BatchUploaderTest.java index bf238a6f29..56d5e70688 100644 --- a/src/test/java/io/stargate/sgv2/jsonapi/service/provider/S3BatchUploaderTest.java +++ b/src/test/java/io/stargate/sgv2/jsonapi/service/billing/S3BatchUploaderTest.java @@ -1,4 +1,4 @@ -package io.stargate.sgv2.jsonapi.service.provider; +package io.stargate.sgv2.jsonapi.service.billing; import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThatThrownBy; @@ -16,6 +16,7 @@ import java.util.UUID; import java.util.concurrent.CompletableFuture; import java.util.regex.Pattern; + import org.junit.jupiter.api.Test; import org.mockito.ArgumentCaptor; import software.amazon.awssdk.core.async.AsyncRequestBody; From 22fa9312b519be321b95ff74bd55e395e289821e Mon Sep 17 00:00:00 2001 From: Aaron Morton Date: Tue, 11 Aug 2026 15:47:00 +1200 Subject: [PATCH 60/65] WIP --- .../jsonapi/config/BillingS3ExportConfig.java | 20 +- .../sgv2/jsonapi/metrics/BillingMetrics.java | 5 +- .../billing/AsyncBatchedLogUploader.java | 33 ++ .../service/billing/BatchedLogBuffer.java | 284 +++++++++++++++ .../sgv2/jsonapi/service/billing/Billing.java | 1 - .../jsonapi/service/billing/BillingQueue.java | 194 ---------- .../billing/BillingS3HandlerInstaller.java | 12 +- .../service/billing/BillingS3LogHandler.java | 334 ++++++++++-------- .../service/billing/DefaultBilling.java | 3 +- .../service/billing/S3BatchUploader.java | 121 ------- .../service/billing/S3BatchedLogUploader.java | 190 ++++++++++ .../service/billing/BillingQueueTest.java | 36 +- .../billing/BillingS3LogHandlerTest.java | 8 +- .../service/billing/S3BatchUploaderTest.java | 16 +- 14 files changed, 734 insertions(+), 523 deletions(-) create mode 100644 src/main/java/io/stargate/sgv2/jsonapi/service/billing/AsyncBatchedLogUploader.java create mode 100644 src/main/java/io/stargate/sgv2/jsonapi/service/billing/BatchedLogBuffer.java delete mode 100644 src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingQueue.java delete mode 100644 src/main/java/io/stargate/sgv2/jsonapi/service/billing/S3BatchUploader.java create mode 100644 src/main/java/io/stargate/sgv2/jsonapi/service/billing/S3BatchedLogUploader.java diff --git a/src/main/java/io/stargate/sgv2/jsonapi/config/BillingS3ExportConfig.java b/src/main/java/io/stargate/sgv2/jsonapi/config/BillingS3ExportConfig.java index 4e89d42c4a..726f4ac584 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/config/BillingS3ExportConfig.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/config/BillingS3ExportConfig.java @@ -19,9 +19,10 @@ public interface BillingS3ExportConfig { /** S3 bucket region */ Optional bucketRegion(); - /** Only for non-AWS S3 endpoints (e.g. S3Mock in tests). - * TODO: XXX EXPLAIN WHAT THIS SHOULD SET SET TO - * */ + /** + * Only for non-AWS S3 endpoints (e.g. S3Mock in tests). TODO: XXX EXPLAIN WHAT THIS SHOULD SET + * SET TO + */ Optional endpointOverride(); /** */ @@ -29,9 +30,9 @@ public interface BillingS3ExportConfig { int maxEventsPerBatch(); /** - * Max bytes to include in a batch, NOTE: if a single event is bigger than this it will be sent in a batch still. - * 2097152 == 2 MB - * */ + * Max bytes to include in a batch, NOTE: if a single event is bigger than this it will be sent in + * a batch still. 2097152 == 2 MB + */ @WithDefault("2097152") long maxBytesPerBatch(); @@ -47,9 +48,10 @@ public interface BillingS3ExportConfig { @WithDefault("4") int uploadConcurrency(); - /** Budget for draining the buffer at shutdown; keep below the pod termination grace period. - * TODO: XXX WHAT IS THE CURRENT TERMINATION PERIOD ? - * */ + /** + * Budget for draining the buffer at shutdown; keep below the pod termination grace period. TODO: + * XXX WHAT IS THE CURRENT TERMINATION PERIOD ? + */ @WithDefault("PT20S") Duration shutdownTimeout(); } diff --git a/src/main/java/io/stargate/sgv2/jsonapi/metrics/BillingMetrics.java b/src/main/java/io/stargate/sgv2/jsonapi/metrics/BillingMetrics.java index a265c809be..a491bfaf58 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/metrics/BillingMetrics.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/metrics/BillingMetrics.java @@ -36,13 +36,10 @@ public final class BillingMetrics { private final Counter batchesFailed; private final AtomicLong lastDeliveryEpochSeconds = new AtomicLong(0); - /** * @param depthSource live queue depth, exposed read-only as {@code billing.s3.queue.depth} */ - public BillingMetrics( - MeterRegistry meterRegistry, Supplier depthSource) { - + public BillingMetrics(MeterRegistry meterRegistry, Supplier depthSource) { this.offered = meterRegistry.counter("billing.s3.events.offered"); this.droppedCapacity = meterRegistry.counter("billing.s3.events.dropped", "reason", "capacity"); diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/billing/AsyncBatchedLogUploader.java b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/AsyncBatchedLogUploader.java new file mode 100644 index 0000000000..dea042c9d0 --- /dev/null +++ b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/AsyncBatchedLogUploader.java @@ -0,0 +1,33 @@ +package io.stargate.sgv2.jsonapi.service.billing; + +import io.smallrye.mutiny.Uni; + +/** + * A function that uploads a batch of log records, normally to S3. + */ +@FunctionalInterface +public interface AsyncBatchedLogUploader extends AutoCloseable { + + /** + * Called to upload the batch of records. + * + * @param batch The batch of log records to upload + * @return A Uni of the result of the operation + */ + Uni upload(BatchedLogBuffer.Batch batch); + + @Override + default void close() {} + + /** + * Result of the upload call. + * @param success true if the operation succeeded, false otherwise. + * @param throwable The throwable associated with an error state. + * @param batch The batch that was uploaded, or attempted to be uploaded. + */ + record UploadResult( + boolean success, + Throwable throwable, + BatchedLogBuffer.Batch batch + ){} +} diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BatchedLogBuffer.java b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BatchedLogBuffer.java new file mode 100644 index 0000000000..cc912fa696 --- /dev/null +++ b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BatchedLogBuffer.java @@ -0,0 +1,284 @@ +package io.stargate.sgv2.jsonapi.service.billing; + +import com.fasterxml.uuid.Generators; +import com.fasterxml.uuid.NoArgGenerator; +import com.google.common.annotations.VisibleForTesting; +import io.stargate.sgv2.jsonapi.metrics.BillingMetrics; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.time.Duration; +import java.time.Instant; +import java.util.*; +import java.util.concurrent.ArrayBlockingQueue; +import java.util.concurrent.BlockingQueue; +import java.util.concurrent.atomic.AtomicLong; +import java.util.logging.LogRecord; + +import static io.stargate.sgv2.jsonapi.util.ClassUtils.classSimpleName; + +/** + * Buffer for {@link LogRecord} that batches them according to the configuration. + *

+ * See {@link #BatchedLogBuffer(int, long, Duration, int, BillingMetrics)} for the config. + *

+ *

+ * There are two uses of this class, producers and consumers. + *

    + *
  • Producers - call {@link #offer(LogRecord)} to add the log record to the buffer.
    • + *
  • Consumers - call {@link #nextBatch(boolean)} to get the next batch to send if there is a + * full batch.
    • + *
+ * + * The buffer is designed to handle these as concurrent calls from different threads, and tracks + * metrics for it's use. + *

+ */ +public class BatchedLogBuffer { + + private static final Logger LOGGER = LoggerFactory.getLogger(BatchedLogBuffer.class); + + private final int maxBatchSize; + private final long maxBytes; + private final Duration maxAge; + + private final BlockingQueue queue; + private final BillingMetrics billingMetrics; + + private final AtomicLong queuedBytes = new AtomicLong(0); + + /** + * Creates a new instance of the buffer. + * + * @param maxBatchSize Maximum number of log records in a batch, when the buffer has more than this many + * entries a new batch is made available which will contain no more than this many lines. + * @param maxBatchBytes Maximum numbers of bytes in a batch, when the buffer has more than this many entries + * a new batch is made available which may contain more than this many bytes. The batch will + * have many maxBatchBytes if there is a single log record that is bigger. + * @param maxAge Maximum age any log record should have in the buffer before a new batch is available. + * @param queueCapacity Total number of log records to buffer. + * @param billingMetrics Metrics recording object. + */ + public BatchedLogBuffer(int maxBatchSize, long maxBatchBytes, Duration maxAge, int queueCapacity, BillingMetrics billingMetrics) { + + if (maxBatchSize < 1) { + throw new IllegalArgumentException("maxBatchSize must be >= 1, got: " + maxBatchSize ); + } + if (maxBatchBytes < 1) { + throw new IllegalArgumentException("maxBatchBytes must be >= 1, got: " + maxBatchBytes ); + } + if (queueCapacity < 1) { + throw new IllegalArgumentException("queueCapacity must be >= 1, got: " + queueCapacity); + } + if (maxAge == null || maxAge.isNegative() || maxAge.isZero()) { + throw new IllegalArgumentException("maxAge must be positive, got: " + maxAge); + } + this.maxBatchSize = maxBatchSize; + this.maxBytes = maxBatchBytes; + this.maxAge = maxAge; + + this.billingMetrics = Objects.requireNonNull(billingMetrics, "billingMetrics must not be null"); + // must be concurrent to handle multiple threads + this.queue = new ArrayBlockingQueue<>(queueCapacity); + } + + /** + * Appends the LogRecord to the buffer if the buffer has capacity. + * + *

+ * NOTE: because this is used for billing information if the record is + * null or has an empty message an exception is thrown rather than + * silently dropping it. We expect this situation to be an exception and it should fail. + *

+ * @param record {@link LogRecord} to add to the buffer. + * @return true if the record was added to be buffer, false if the buffer did not have capacity. NOTE: + * this is different to the param check for record, the buffer filling is tracked as metric but no error. + * + */ + public boolean offer(LogRecord record) { + + Objects.requireNonNull(record, "record must not be null"); + + var logLine = record.getMessage(); + if (logLine == null || logLine.isBlank()) { + throw new IllegalArgumentException("record.getMessage() must not be null or blank"); + } + var newEntry = new Entry(record.getInstant(), logLine); + + billingMetrics.recordOffered(); + if (!queue.offer(newEntry)) { + // Bounded buffer full: drop and count + billingMetrics.recordDropped(); + return false; + } + + queuedBytes.addAndGet(newEntry.lineBytes()); + return true; + } + + /** + * Returns the next batch of messages from the {@link LogRecord}'s added to the buffer, + * if one is available. + *

+ * Designed to be called from different threads than those producing LogRecord's. + *

+ * + * @param drainFully when True a new batch is created without checking the configured + * rules, use this when draining the buffer and there may only + * be a partial batch. + * @return A new {@link Batch} of log messages all of which have been removed from the + * buffer, or null if there is no next batch. + * */ + public Batch nextBatch(boolean drainFully) { + + var batchReason = decideNextBatch(drainFully); + if (batchReason == null) { + return null; + } + + List lines = new ArrayList<>(maxBatchSize); + Instant oldestEventAt = null; + long batchBytes = 0; + Entry entry; + + // No matter why we started we create a full batch, e.g. we could start because the oldest + // entry is past maxAge, but we still fill the batch. + while (lines.size() < maxBatchSize && batchBytes < maxBytes && (entry = queue.poll()) != null) { + + if (oldestEventAt == null || entry.eventAt().isBefore(oldestEventAt)) { + oldestEventAt = entry.eventAt(); + } + + lines.add(entry.line()); + var lineBytes = entry.lineBytes(); + queuedBytes.addAndGet(-lineBytes); + batchBytes += lineBytes; + } + + // sanity check, in case of concurrent calls + if (lines.isEmpty()){ + return null; + } + + return new Batch(batchReason, lines, oldestEventAt); + } + + public boolean isEmpty() { + return queue.isEmpty(); + } + + public int size() { + return queue.size(); + } + + @VisibleForTesting + public long queuedBytes() { + return queuedBytes.get(); + } + + public int remainingCapacity() { + return queue.remainingCapacity(); + } + + + + private Duration oldestEntry() { + var head = queue.peek(); + + return head == null ? Duration.ZERO : Duration.between(head.eventAt(), Instant.now()); + } + + + + + private BillingBatchReason decideNextBatch(boolean drainFully) { + + if (drainFully){ + return BillingBatchReason.DRAINING; + } + if (queue.size() >= maxBatchSize) { + return BillingBatchReason.MAX_BATCH_SIZE_EXCEEDED + } + if (queuedBytes.get() > maxBytes) { + return BillingBatchReason.MAX_BYTES_EXCEEDED; + } + if (oldestEntry().compareTo(maxAge) >= 0){ + return BillingBatchReason.MAX_AGE_EXCEEDED; + } + return null; + } + + public enum BillingBatchReason { + DRAINING, + MAX_BATCH_SIZE_EXCEEDED, + MAX_BYTES_EXCEEDED, + MAX_AGE_EXCEEDED + } + + /** + * OLD BELOW + * + * One drained, sealed batch. {@code oldestEventAt} is the minimum event time across {@code lines} + * — queue order is enqueue order, not event-time order, under concurrent publish. + */ + public static final class Batch { + + private static final NoArgGenerator UUID_V7_GENERATOR = Generators.timeBasedEpochGenerator(); + + + private final UUID batchId = UUID_V7_GENERATOR.generate(); + private final BillingBatchReason batchReason; + private final List lines; + private final Instant oldestEventAt; + + public Batch(BillingBatchReason batchReason, List lines, Instant oldestEventAt) { + this.batchReason = batchReason; + this.lines = Collections.unmodifiableList(lines); + this.oldestEventAt = oldestEventAt; + } + + public UUID id() { + return batchId; + } + + public BillingBatchReason reason() { + return batchReason; + } + + public List lines() { + return lines; + } + + public Instant oldestEventAt() { + return oldestEventAt; + } + + public int size() { + return lines.size(); + } + + public String description() { + return "id:%s, reason:%s, oldestEventAt:%s, size:%s".formatted( + batchId, batchReason, oldestEventAt, size() + ); + } + } + + /** + * Holder for the billing event lines we get called with. + * @param eventAt When the event happened + * @param line The billing event line to record + */ + private record Entry(Instant eventAt, String line) { + + /** + * Gets the length of the line in bytes, + *

+ * Kind of a hack, we are counting unicode code points and calling that 1 byte. Should work + * for ascii text, will undercount if there is non ascii chars but everthing in billing should be ascii + */ + public int lineBytes() { + return line.length() + 1; // +1 is for a newline + } + } +} diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/billing/Billing.java b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/Billing.java index 3b2a92e2a6..a0c3fd50a8 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/service/billing/Billing.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/Billing.java @@ -4,7 +4,6 @@ import io.stargate.sgv2.jsonapi.config.feature.ApiFeature; import io.stargate.sgv2.jsonapi.config.feature.ApiFeatures; import io.stargate.sgv2.jsonapi.service.provider.ModelUsage; - import java.util.Objects; /** diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingQueue.java b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingQueue.java deleted file mode 100644 index a34b118712..0000000000 --- a/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingQueue.java +++ /dev/null @@ -1,194 +0,0 @@ -package io.stargate.sgv2.jsonapi.service.billing; - -import com.google.common.annotations.VisibleForTesting; - -import java.time.Duration; -import java.time.Instant; -import java.util.ArrayList; -import java.util.Collections; -import java.util.List; -import java.util.concurrent.ArrayBlockingQueue; -import java.util.concurrent.BlockingQueue; -import java.util.concurrent.atomic.AtomicLong; - -/** - * Bounded in-memory buffer that owns the batching policy of the billing S3 export: it decides when - * a batch is sealed ({@code maxEvents} lines or {@code maxBytes} UTF-8 NDJSON bytes) and hands out - * drained {@link Batch}es. A batch may exceed {@code maxBytes} by one whole line; lines are never - * split. - */ -public class BillingQueue { - - private final BlockingQueue queue; - - // Approximate buffered bytes (see lineBytes); - private final AtomicLong queuedBytes = new AtomicLong(0); - private final int maxBatchSize; - private final long maxBytes; - private final Duration maxAge; - - public BillingQueue(int maxBatchSize, long maxBatchBytes, Duration maxAge, int queueCapacity) { - - if (maxBatchSize < 1) { - throw new IllegalArgumentException("maxBatchSize must be >= 1, got: " + maxBatchSize ); - } - if (maxBatchBytes < 1) { - throw new IllegalArgumentException("maxBatchBytes must be >= 1, got: " + maxBatchBytes ); - } - if (queueCapacity < 1) { - throw new IllegalArgumentException("queueCapacity must be >= 1, got: " + queueCapacity); - } - if (maxAge == null || maxAge.isNegative() || maxAge.isZero()) { - throw new IllegalArgumentException("maxAge must be positive, got: " + maxAge); - } - this.maxBatchSize = maxBatchSize; - this.maxBytes = maxBatchBytes; - this.maxAge = maxAge; - - // must be concurrent to handle multiple threads - this.queue = new ArrayBlockingQueue<>(queueCapacity); - } - - public boolean isEmpty() { - return queue.isEmpty(); - } - - public int size() { - return queue.size(); - } - - /** Buffered-bytes counter, exposed for accounting assertions only (no production caller). */ - @VisibleForTesting - long queuedBytes() { - return queuedBytes.get(); - } - - /** - * Appends a line to the billing queue - * - * @param eventAt when the event was logged - * @param line String for the JSON billing blob - * @return true if the billing line was added to the quee for publishing, false otherwise which - * means the queue has reached capacity and we can no longer buffer events - * - * // AI BELOW - * Buffers one line, or returns {@code false} when the capacity bound is hit. - * - * @param eventAt when the event was logged; carried through to {@link Batch#oldestEventAt()} - */ - public boolean offer(Instant eventAt, String line) { - - var newEntry = new Entry(eventAt, line); - if (!queue.offer(newEntry)) { - return false; - } - queuedBytes.addAndGet(newEntry.lineBytes()); - return true; - } - - private Duration oldestEntry() { - var head = queue.peek(); - - return head == null ? Duration.ZERO : Duration.between(head.eventAt(), Instant.now()); - } - - /** - * Tests if we should start a new batch - * */ - private BillingBatchReason startNextBatch() { - - if (queue.size() >= maxBatchSize) { - return BillingBatchReason.MAX_BATCH_SIZE_EXCEEDED - } - if (queuedBytes.get() > maxBytes) { - return BillingBatchReason.MAX_BYTES_EXCEEDED; - } - if (oldestEntry().compareTo(maxAge) >= 0){ - return BillingBatchReason.MAX_AGE_EXCEEDED; - } - return null; - } - - /** - * - * Assumes we are not running async - * OLD BELOW: - * - * Removes and returns up to one sealed batch - * */ - public Batch maybeDrain() { - - var batchReason = startNextBatch(); - if (batchReason == null) { - return null; - } - - List lines = new ArrayList<>(maxBatchSize); - Instant oldestEventAt = null; - long batchBytes = 0; - Entry entry; - - // No matter why we started we create a full batch, e.g. we could start because the oldest - // entry is past maxAge, but we still fill the batch. - while (lines.size() < maxBatchSize && batchBytes < maxBytes && (entry = queue.poll()) != null) { - - if (oldestEventAt == null || entry.eventAt().isBefore(oldestEventAt)) { - oldestEventAt = entry.eventAt(); - } - - lines.add(entry.line()); - var lineBytes = entry.lineBytes(); - queuedBytes.addAndGet(-lineBytes); - batchBytes += lineBytes; - } - - // sanity check, in case of concurrent calls - if (lines.isEmpty()){ - return null; - } - - return new Batch(batchReason, lines, oldestEventAt); - } - - - public enum BillingBatchReason { - MAX_BATCH_SIZE_EXCEEDED, - MAX_BYTES_EXCEEDED, - MAX_AGE_EXCEEDED - } - - /** - * OLD BELOW - * - * One drained, sealed batch. {@code oldestEventAt} is the minimum event time across {@code lines} - * — queue order is enqueue order, not event-time order, under concurrent publish. - */ - public record Batch(BillingBatchReason batchReason, List lines, Instant oldestEventAt) { - - public Batch{ - lines = Collections.unmodifiableList(lines); - } - - public int size(){ - return lines.size(); - } - } - - /** - * Holder for the billing event lines we get called with. - * @param eventAt When the event happened - * @param line The billing event line to record - */ - private record Entry(Instant eventAt, String line) { - - /** - * Gets the length of the line in bytes, - *

- * Kind of a hack, we are counting unicode code points and calling that 1 byte. Should work - * for ascii text, will undercount if there is non ascii chars but everthing in billing should be ascii - */ - public int lineBytes() { - return line.length() + 1; // +1 is for a newline - } - } -} diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3HandlerInstaller.java b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3HandlerInstaller.java index 5c57548f17..0d4a39109c 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3HandlerInstaller.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3HandlerInstaller.java @@ -11,12 +11,11 @@ import org.slf4j.LoggerFactory; /** - * * TODO: XXX MAKE THIS COMMENTS READABLE BY A HUMAN * - * Attaches a {@link BillingS3LogHandler} to the {@code billing.events} JUL logger at startup (when - * {@link BillingS3ExportConfig#enabled()} is {@code true}) and removes + closes it on shutdown for - * a graceful drain. + *

Attaches a {@link BillingS3LogHandler} to the {@code billing.events} JUL logger at startup + * (when {@link BillingS3ExportConfig#enabled()} is {@code true}) and removes + closes it on + * shutdown for a graceful drain. * *

Done programmatically because Quarkus config can't express it: a category's {@code handlers} * list can only reference Quarkus's built-in handler types (console/file/syslog/socket), not a @@ -29,7 +28,8 @@ @ApplicationScoped public class BillingS3HandlerInstaller { - private static final org.slf4j.Logger LOGGER = LoggerFactory.getLogger(BillingS3HandlerInstaller.class); + private static final org.slf4j.Logger LOGGER = + LoggerFactory.getLogger(BillingS3HandlerInstaller.class); static final String BILLING_LOGGER_NAME = "billing.events"; @@ -55,7 +55,7 @@ void onStart(@Observes StartupEvent event) { var bucket = config.bucket().orElse(null); // Fail-loud: invalid billing S3 config throws here, aborting application startup. - var uploader = S3BatchUploader.create(region, bucket, config.endpointOverride()); + var uploader = S3BatchedLogUploader.create(region, bucket, config.endpointOverride()); this.handler = new BillingS3LogHandler(config, uploader, meterRegistry); // TODO: LOGGER NAME SHOULD BE IN CONFIG diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3LogHandler.java b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3LogHandler.java index 88cc82f123..ff900c6f26 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3LogHandler.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3LogHandler.java @@ -2,7 +2,7 @@ import com.google.common.annotations.VisibleForTesting; import io.smallrye.mutiny.Uni; - +import io.stargate.sgv2.jsonapi.metrics.BillingMetrics; import java.time.Duration; import java.util.Objects; import java.util.concurrent.CountDownLatch; @@ -10,22 +10,17 @@ import java.util.concurrent.atomic.AtomicBoolean; import java.util.logging.Handler; import java.util.logging.LogRecord; - -import io.stargate.sgv2.jsonapi.metrics.BillingMetrics; import org.slf4j.Logger; import org.slf4j.LoggerFactory; /** * A Logging handler designed to be used wioth the Billing system. It accpets billing event log * messges, batches them, and then sends to S3. - *

- * See {@link BillingS3HandlerInstaller} for setup. - *

* - * // AI SLOP BELOW - * JUL handler that turns {@code billing.events} log lines into batched S3 objects. + *

See {@link BillingS3HandlerInstaller} for setup. // AI SLOP BELOW JUL handler that turns + * {@code billing.events} log lines into batched S3 objects. * - *

Division of labor: {@link BillingQueue} decides when a batch seals, {@link AsyncBatchUploader} + *

Division of labor: {@link BatchedLogBuffer} decides when a batch seals, {@link AsyncBatchedLogUploader} * decides what an S3 object looks like, and this class decides when uploads run — the flush * triggers (seal on publish, age tick, drain on close), the upload-concurrency gate, and metrics. * @@ -34,24 +29,43 @@ */ public final class BillingS3LogHandler extends Handler { + // Logger for this handler, not the destination we are sending events to. private static final Logger LOGGER = LoggerFactory.getLogger(BillingS3LogHandler.class); + /** + * Duration the upload thread sleeps between uploading. After all available batches have been + * uploaded, sleeps for this long waiting for the {@link #wakeupSignal} + */ + private static final long UPLOAD_SLEEP_MS = 1000; + + /** + * When signaled this object wakes up the uploading thread to immediately get to work. Used as + * part of the close mechanism to trigger uploading to complete. We do not signal the upload + * everytime a producer calls {@link #publish(LogRecord)}. + */ private final Object wakeupSignal = new Object(); + + /** + * When true means the Handler has been closed via {@link #close()} and it will silently drop any + * further calls to publish log entries. This also cauese the upload thread to empty the queue + */ private final AtomicBoolean isClosed = new AtomicBoolean(false); - private final AtomicBoolean isDraining = new AtomicBoolean(false); - private final CountDownLatch finishedLatch = new CountDownLatch(1); - private final AsyncBatchUploader uploader; + /** + * Started at 1 and then decremented in {@link #startUploading()} when it exists so we know we + * have finished uploading. + */ + private final CountDownLatch uploadingFinished = new CountDownLatch(0); + + private final AsyncBatchedLogUploader uploader; private final BillingMetrics billingMetrics; - private final BillingQueue billingQueue; + private final BatchedLogBuffer batchedLogBuffer; @VisibleForTesting BillingS3LogHandler( - AsyncBatchUploader uploader, - BillingQueue billingQueue, - BillingMetrics billingMetrics) { + AsyncBatchedLogUploader uploader, BatchedLogBuffer batchedLogBuffer, BillingMetrics billingMetrics) { - this.billingQueue = billingQueue; + this.batchedLogBuffer = batchedLogBuffer; this.uploader = uploader; this.billingMetrics = Objects.requireNonNull(billingMetrics); } @@ -71,33 +85,29 @@ private static Duration requirePositive(Duration value, String property) { @Override public void publish(LogRecord record) { - if (record == null || isClosed.get()) { + // Sanity check + if (record == null) { return; } - // TODO: XXX : WHAT DOES MEAN "This handler never runs a Formatter, so a parameterized call would ship its" - - // Producer contract (DefaultBilling): the message is the final JSON line, logged without {} - // placeholders. This handler never runs a Formatter, so a parameterized call would ship its - // raw template. getInstant() is when the producer logged it — within microseconds of the - // "timestamp" it embedded in the JSON, and the object key only needs minute resolution. - - var line = record.getMessage(); - if (line == null || line.isBlank()) { - return; + if (isClosed.get()) { + LOGGER.warn("publish() - called when closed, dropping record:{}", record); } - billingMetrics.recordOffered(); - if (!billingQueue.offer(record.getInstant(), line)) { - // Bounded buffer full: drop and count - billingMetrics.recordDropped(); + // buffer handles metrics + if (!batchedLogBuffer.offer(record)) { + if (LOGGER.isTraceEnabled()) { + LOGGER.trace("publish() - dropped record:{}", record); + } } // flushing runs every second so nothing more to do } @Override public void flush() { - setWakeupSignal(); + // wakeup the upload thread to send eveything it can. + // NOTE: this will only drain the buffer if isClosed() is true + notifyUploading(); } /** @@ -108,43 +118,41 @@ public void flush() { @Override public void close() { + // mark as closed to stop accepting further events and tell the upload thread + // to drain the bugger/ isClosed.set(true); - isDraining.set(true); flush(); try { - if (!finishedLatch.await(30, TimeUnit.SECONDS)) { - LOGGER.warn("close() - Billing upload loop did not stop within 30s, interrupting"); - } + // waiting for the uploading thead to signal it has sent all events in the buffer + if (!uploadingFinished.await(30, TimeUnit.SECONDS)) { + LOGGER.warn("close() - Billing upload loop did not stop within 30s, interrupting"); + } } catch (InterruptedException e) { Thread.currentThread().interrupt(); LOGGER.warn("close() - Interrupted waiting for billing upload loop to finish"); - } - finally { + } finally { uploader.close(); } - } // ============================================================ // Flush pipeline // ============================================================ - private void setWakeupSignal() { - synchronized (wakeupSignal) { - wakeupSignal.notifyAll(); - } - } - - void startPublishing() { + /** Called on a worker thread to start uploading log records. */ + void startUploading() { - BillingQueue.Batch batch; + BatchedLogBuffer.Batch batch; try { while (true) { synchronized (wakeupSignal) { - if (!isDraining.get()) { + // if the handler is closed we do not want to go to sleep again because it is closing + // down. + if (!isClosed.get()) { try { + // waiting will release the synchronized monitor wakeupSignal.wait(1000); } catch (InterruptedException e) { Thread.currentThread().interrupt(); @@ -153,132 +161,146 @@ void startPublishing() { } } - while ((batch = billingQueue.maybeDrain()) != null) { + // Get the next batches, if isClosed is true then we want to drain all events + // which may mean creating a batch when we do not have a full one. + + while ((batch = batchedLogBuffer.nextBatch(isClosed.get())) != null) { + + // calling await() on the Uni from deferBatch causes the uni to start running // indefinitely() is bounded by the 10s ifNoItem() timeout inside deferBatch(), // and failures are recovered there, so this neither hangs nor throws. deferBatch(batch).await().indefinitely(); } - if (isDraining.get()){ - return; + if (isClosed.get()) { + // Handler is closing down, time to get out of this crazy loop + break; } } - } - finally { - billingMetrics.recordAbandonedAtShutdown(billingQueue.size()); - if (!billingQueue.isEmpty()){ - LOGGER.warn("start() - finished with abandoned billing events, billingQueue.size():{} " , billingQueue.size()); + } finally { + // record if there are any abandonded events + billingMetrics.recordAbandonedAtShutdown(batchedLogBuffer.size()); + if (!batchedLogBuffer.isEmpty()) { + LOGGER.warn( + "start() - finished with abandoned billing events, billingQueue.size():{} ", + batchedLogBuffer.size()); } - finishedLatch.countDown(); + + // reset the latch so the call at close() can exit. + uploadingFinished.countDown(); + } + } + + /** + * Signals to the uploading thread that it should wakeup and do some work. + * + *

The uploading thread runs repeated checks for new batches, this is only needed to wakeup as + * part of closing + */ + private void notifyUploading() { + synchronized (wakeupSignal) { + wakeupSignal.notifyAll(); } } - private Uni deferBatch(BillingQueue.Batch batch) { + /** + * Creates a Uni that will upload the provided batch. + * + *

As a deferred Uni it does not do any work until something pulls the item, so the caller (see + * startUploading()) starts the work and can decide to wait etc. + * + * @param batch + * @return + */ + private Uni deferBatch(BatchedLogBuffer.Batch batch) { // upload() is called at subscription, not when this method returns. // deferred also converts a synchronous throw from upload() into a Uni failure. return Uni.createFrom() - .deferred(() -> uploader.upload(batch)) - .ifNoItem() - .after(Duration.ofSeconds(10)) - .fail() - .onItemOrFailure() - .invoke( - (item, failure) -> { - if (failure != null) { - billingMetrics.recordBatchFailed(batch.size()); - LOGGER.error("Failed to upload billing S3 batch ({} events)", batch.size(), failure); - } else { - billingMetrics.recordBatchDelivered(batch.size()); - } - }) - .onFailure() - .recoverWithNull(); + .deferred(() -> uploader.upload(batch)) + .ifNoItem() + .after(Duration.ofSeconds(10)) + .fail(); } -// /** Seal-triggered flush: ship when the buffer has a full batch by count or bytes. */ -// private void maybeFlush() { -// if (eventQueue.shouldFlush()) { -// tryFlush(); -// } -// } - -// /** -// * Age trigger: every {@code maxAge} tick ships whatever is buffered, sealed or not. Deliberately -// * no head-age check: flushing only entries older than {@code maxAge} would let an event that just -// * missed a tick wait ~2x{@code maxAge}, while shipping unconditionally bounds every wait by one -// * period — at the cost of an occasional small object when a tick lands just after a seal flush. -// * -// *

Catches everything: an escaped throwable would silently cancel all future runs of a -// * fixed-rate task. -// */ -// @VisibleForTesting -// void onAgeTick() { -// try { -// if (!eventQueue.isEmpty()) { -// tryFlush(); -// } -// } catch (Throwable t) { -// LOG.error("Billing S3 export age-flush tick failed", t); -// } -// } -// -// /** -// * Claims an in-flight slot (non-blocking CAS, at most {@link #uploadConcurrency} held) and, on -// * success, drains + uploads one batch asynchronously. When the upload settles the slot is -// * released and the seal condition re-checked: a full batch may have accumulated meanwhile. -// */ -// private void tryFlush() { -// -// int prev = inFlight.getAndUpdate(n -> n < uploadConcurrency ? n + 1 : n); -// if (prev >= uploadConcurrency) { -// return; -// } -// Uni.createFrom() -// .item(eventQueue::maybeDrain) -// .runSubscriptionOn(Infrastructure.getDefaultWorkerPool()) -// .flatMap(this::uploadBatch) -// .eventually( -// () -> { -// inFlight.getAndDecrement(); -// maybeFlush(); -// }) -// .subscribe() -// .with(ignored -> {}, failure -> LOG.error("Billing S3 export flush failed", failure)); -// } + // /** Seal-triggered flush: ship when the buffer has a full batch by count or bytes. */ + // private void maybeFlush() { + // if (eventQueue.shouldFlush()) { + // tryFlush(); + // } + // } + + // /** + // * Age trigger: every {@code maxAge} tick ships whatever is buffered, sealed or not. + // Deliberately + // * no head-age check: flushing only entries older than {@code maxAge} would let an event that + // just + // * missed a tick wait ~2x{@code maxAge}, while shipping unconditionally bounds every wait by + // one + // * period — at the cost of an occasional small object when a tick lands just after a seal + // flush. + // * + // *

Catches everything: an escaped throwable would silently cancel all future runs of a + // * fixed-rate task. + // */ + // @VisibleForTesting + // void onAgeTick() { + // try { + // if (!eventQueue.isEmpty()) { + // tryFlush(); + // } + // } catch (Throwable t) { + // LOG.error("Billing S3 export age-flush tick failed", t); + // } + // } + // + // /** + // * Claims an in-flight slot (non-blocking CAS, at most {@link #uploadConcurrency} held) and, + // on + // * success, drains + uploads one batch asynchronously. When the upload settles the slot is + // * released and the seal condition re-checked: a full batch may have accumulated meanwhile. + // */ + // private void tryFlush() { + // + // int prev = inFlight.getAndUpdate(n -> n < uploadConcurrency ? n + 1 : n); + // if (prev >= uploadConcurrency) { + // return; + // } + // Uni.createFrom() + // .item(eventQueue::maybeDrain) + // .runSubscriptionOn(Infrastructure.getDefaultWorkerPool()) + // .flatMap(this::uploadBatch) + // .eventually( + // () -> { + // inFlight.getAndDecrement(); + // maybeFlush(); + // }) + // .subscribe() + // .with(ignored -> {}, failure -> LOG.error("Billing S3 export flush failed", failure)); + // } /** Uploads one batch; never fails the pipeline — a batch that exhausts retries is counted. */ -// private Uni uploadBatch(BillingQueue.Batch batch) { -// if (batch.isEmpty()) { -// return Uni.createFrom().voidItem(); -// } -// int size = batch.size(); -// // Runs immediately on subscription; deferred only turns a throw before upload() returns a Uni -// // into a Uni failure handled below. -// return Uni.createFrom() -// .deferred(() -> uploader.upload(batch)) -// .onItem() -// .invoke(() -> billingMetrics.recordBatchDelivered(size)) -// .onFailure() -// .invoke(t -> LOG.error("Failed to upload billing S3 batch ({} events)", size, t)) -// .onFailure() -// .recoverWithItem( -// () -> { -// billingMetrics.recordBatchFailed(size); -// return null; -// }); -// } + // private Uni uploadBatch(BillingQueue.Batch batch) { + // if (batch.isEmpty()) { + // return Uni.createFrom().voidItem(); + // } + // int size = batch.size(); + // // Runs immediately on subscription; deferred only turns a throw before upload() returns a + // Uni + // // into a Uni failure handled below. + // return Uni.createFrom() + // .deferred(() -> uploader.upload(batch)) + // .onItem() + // .invoke(() -> billingMetrics.recordBatchDelivered(size)) + // .onFailure() + // .invoke(t -> LOG.error("Failed to upload billing S3 batch ({} events)", size, t)) + // .onFailure() + // .recoverWithItem( + // () -> { + // billingMetrics.recordBatchFailed(size); + // return null; + // }); + // } - /** - * Uploads one sealed batch to the export destination; owns the object key and body encoding. - * Implementations must tolerate concurrent calls. - */ - @FunctionalInterface - public interface AsyncBatchUploader extends AutoCloseable { - Uni upload(BillingQueue.Batch batch); - - @Override - default void close() {} - } } diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/billing/DefaultBilling.java b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/DefaultBilling.java index f149226ab3..b62e8c5ffb 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/service/billing/DefaultBilling.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/DefaultBilling.java @@ -8,14 +8,13 @@ import com.google.common.annotations.VisibleForTesting; import io.stargate.sgv2.jsonapi.config.BillingConfig; import io.stargate.sgv2.jsonapi.config.feature.ApiFeature; +import io.stargate.sgv2.jsonapi.service.provider.ModelUsage; import java.time.Instant; import java.util.ArrayList; import java.util.List; import java.util.Objects; import java.util.Set; import java.util.UUID; - -import io.stargate.sgv2.jsonapi.service.provider.ModelUsage; import org.slf4j.Logger; import org.slf4j.LoggerFactory; diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/billing/S3BatchUploader.java b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/S3BatchUploader.java deleted file mode 100644 index 8cf1758086..0000000000 --- a/src/main/java/io/stargate/sgv2/jsonapi/service/billing/S3BatchUploader.java +++ /dev/null @@ -1,121 +0,0 @@ -package io.stargate.sgv2.jsonapi.service.billing; - -import com.google.common.annotations.VisibleForTesting; -import io.smallrye.mutiny.Uni; -import java.net.URI; -import java.nio.charset.StandardCharsets; -import java.time.Duration; -import java.time.Instant; -import java.time.ZoneOffset; -import java.time.format.DateTimeFormatter; -import java.util.List; -import java.util.Objects; -import java.util.Optional; -import java.util.UUID; - -import software.amazon.awssdk.core.async.AsyncRequestBody; -import software.amazon.awssdk.regions.Region; -import software.amazon.awssdk.services.s3.S3AsyncClient; -import software.amazon.awssdk.services.s3.model.PutObjectRequest; - -/** Uploads sealed billing batches to S3 as NDJSON objects under time-partitioned keys. */ -public class S3BatchUploader implements BillingS3LogHandler.AsyncBatchUploader { - - // S3 object-key consistent identifier - static final String PATH_PREFIX = "data-api"; - private static final String NDJSON_CONTENT_TYPE = "application/x-ndjson"; - // object key format - private static final DateTimeFormatter KEY_TIME_FORMAT = - DateTimeFormatter.ofPattern("yyyy/MM/dd/HH/mm").withZone(ZoneOffset.UTC); - - // Bound every PUT so a hung connection can neither pin an upload slot indefinitely nor stall - // the shutdown drain. Retries stay inside the SDK's built-in default policy (bounded attempts, - // jittered throttle-aware backoff). - private static final Duration API_CALL_ATTEMPT_TIMEOUT = Duration.ofSeconds(10); - private static final Duration API_CALL_TIMEOUT = Duration.ofSeconds(30); - - private final S3AsyncClient client; - private final String bucket; - - S3BatchUploader(S3AsyncClient client, String bucket) { - this.client = client; - this.bucket = bucket; - } - - public static S3BatchUploader create( - - String region, String bucket, Optional endpointOverride) { - if (region == null || region.isBlank()) { - throw new IllegalArgumentException("stargate.jsonapi.billing.s3.bucket-region must be set"); - } - - if (bucket == null || bucket.isBlank()) { - throw new IllegalArgumentException("stargate.jsonapi.billing.s3.bucket must be set"); - } - Objects.requireNonNull(endpointOverride, "endpointOverride must not be null"); - - // Credentials resolve from the SDK's default provider chain (env vars, web-identity/OIDC - // token, instance/container roles), left implicit so the client owns — and closes — the - // provider. This transparently supports federated (AssumeRoleWithWebIdentity) and - // cross-account access: the bucket may live in a different account (per IAM + bucket - // policy); its region is set via .region(). - var builder = - S3AsyncClient.builder() - .region(Region.of(region)) - .overrideConfiguration( - o -> - o.apiCallAttemptTimeout(API_CALL_ATTEMPT_TIMEOUT) - .apiCallTimeout(API_CALL_TIMEOUT)); - - // Real AWS S3 needs no endpoint: the SDK endpoint rules (s3 SDK's DefaultS3EndpointProvider) - // derive https://.s3..amazonaws.com from region + partition dnsSuffix. - // An override is only for a non-AWS S3 (S3Mock in tests): it bypasses those rules and forces - // path-style, since a localhost host can't virtual-host the bucket as a subdomain. - endpointOverride - .filter(s -> !s.isBlank()) - .ifPresent(uri -> builder.endpointOverride(URI.create(uri)).forcePathStyle(true)); - - return new S3BatchUploader(builder.build(), bucket); - } - - @Override - public Uni upload(BillingQueue.Batch batch) { - String key = objectKey(batch.oldestEventAt(), UUID.randomUUID()); - byte[] body = toNdjson(batch.lines()); - // No .retry() here: unconfigured, S3AsyncClient already retries (default LegacyRetryStrategy — - // 3 retries / 4 attempts). See - // https://docs.aws.amazon.com/sdk-for-java/latest/developer-guide/retry-strategy.html - // and see https://github.com/aws/aws-sdk-java-v2/issues/6987 for future change. - return Uni.createFrom() - .completionStage( - () -> - client - .putObject( - PutObjectRequest.builder() - .bucket(bucket) - .key(key) - .contentType(NDJSON_CONTENT_TYPE) - .build(), - AsyncRequestBody.fromBytes(body)) - .thenAccept(resp -> {})); - } - - @VisibleForTesting - static String objectKey(Instant timestamp, UUID id) { - return PATH_PREFIX + "/" + KEY_TIME_FORMAT.format(timestamp) + "/" + id + ".jsonl"; - } - - @VisibleForTesting - static byte[] toNdjson(List lines) { - StringBuilder sb = new StringBuilder(); - for (String line : lines) { - sb.append(line).append('\n'); - } - return sb.toString().getBytes(StandardCharsets.UTF_8); - } - - @Override - public void close() { - client.close(); - } -} diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/billing/S3BatchedLogUploader.java b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/S3BatchedLogUploader.java new file mode 100644 index 0000000000..72a9bcef7d --- /dev/null +++ b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/S3BatchedLogUploader.java @@ -0,0 +1,190 @@ +package io.stargate.sgv2.jsonapi.service.billing; + +import com.google.common.annotations.VisibleForTesting; +import io.smallrye.mutiny.Uni; +import java.net.URI; +import java.nio.charset.StandardCharsets; +import java.time.Duration; +import java.time.ZoneOffset; +import java.time.format.DateTimeFormatter; +import java.util.Objects; +import java.util.concurrent.CompletionException; + +import io.stargate.sgv2.jsonapi.metrics.BillingMetrics; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import software.amazon.awssdk.awscore.exception.AwsServiceException; +import software.amazon.awssdk.core.async.AsyncRequestBody; +import software.amazon.awssdk.regions.Region; +import software.amazon.awssdk.services.s3.S3AsyncClient; +import software.amazon.awssdk.services.s3.model.PutObjectRequest; + +/** Uploads sealed billing batches to S3 as NDJSON objects under time-partitioned keys. */ +public class S3BatchedLogUploader implements AsyncBatchedLogUploader { + + private static final Logger LOGGER = LoggerFactory.getLogger(S3BatchedLogUploader.class); + + // S3 destination formatting + private static final String PATH_PREFIX = "data-api"; + private static final String CONTENT_TYPE_NDJSON = "application/x-ndjson"; + private static final DateTimeFormatter OBJECT_KEY_FORMAT = + DateTimeFormatter.ofPattern("yyyy/MM/dd/HH/mm").withZone(ZoneOffset.UTC); + + // Bound every PUT so a hung connection can neither pin an upload slot indefinitely nor stall + // the shutdown drain. Retries stay inside the SDK's built-in default policy (bounded attempts, + // jittered throttle-aware backoff). + private static final Duration API_CALL_ATTEMPT_TIMEOUT = Duration.ofSeconds(10); + private static final Duration API_CALL_TIMEOUT = Duration.ofSeconds(30); + + private final S3AsyncClient client; + private final String bucket; + private final BillingMetrics billingMetrics; + + @VisibleForTesting + S3BatchedLogUploader(S3AsyncClient client, String bucket, BillingMetrics billingMetrics) { + this.client = client; + this.bucket = bucket; + this.billingMetrics = billingMetrics; + } + + /** + * Creates a new instance + * + * @param region + * @param bucket + * @param endpointOverride + * @return + */ + public static S3BatchedLogUploader create(String region, String bucket, String endpointOverride, BillingMetrics billingMetrics) { + + if (region == null || region.isBlank()) { + throw new IllegalArgumentException("region must be set"); + } + if (bucket == null || bucket.isBlank()) { + throw new IllegalArgumentException("bucket must be set"); + } + + // Credentials resolve from the SDK's default provider chain (env vars, web-identity/OIDC + // token, instance/container roles), left implicit so the client owns — and closes — the + // provider. This transparently supports federated (AssumeRoleWithWebIdentity) and + // cross-account access: the bucket may live in a different account (per IAM + bucket + // policy); its region is set via .region(). + var builder = + S3AsyncClient.builder() + .region(Region.of(region)) + .overrideConfiguration( + config -> + config + .apiCallAttemptTimeout(API_CALL_ATTEMPT_TIMEOUT) + .apiCallTimeout(API_CALL_TIMEOUT)); + + // Real AWS S3 needs no endpoint: the SDK endpoint rules (s3 SDK's DefaultS3EndpointProvider) + // derive https://.s3..amazonaws.com from region + partition dnsSuffix. + // An override is only for a non-AWS S3 (S3Mock in tests): it bypasses those rules and forces + // path-style, since a localhost host can't virtual-host the bucket as a subdomain. + if (endpointOverride != null) { + builder.endpointOverride(URI.create(endpointOverride)).forcePathStyle(true); + } + + return new S3BatchedLogUploader(builder.build(), bucket); + } + + @Override + public Uni upload(BatchedLogBuffer.Batch batch) { + + Objects.requireNonNull(batch, "batch must not be null"); + + var key = objectKey(batch); + var body = objectContent(batch); + + if (LOGGER.isDebugEnabled()) { + LOGGER.debug( + "upload() - starting to upload batch, batch:({}), S3.bucket:{}, S3.key:{}", + batch.description(), + bucket, + key); + } + + // No .retry() here: unconfigured, S3AsyncClient already retries (default LegacyRetryStrategy — + // 3 retries / 4 attempts). See + // https://docs.aws.amazon.com/sdk-for-java/latest/developer-guide/retry-strategy.html + // and see https://github.com/aws/aws-sdk-java-v2/issues/6987 for future change. + + return Uni.createFrom() + .completionStage( + () -> + client.putObject( + PutObjectRequest.builder() + .bucket(bucket) + .key(key) + .contentType(CONTENT_TYPE_NDJSON) + .build(), + AsyncRequestBody.fromBytes(body))) + .onItemOrFailure() + .transform( + (resp, failure) -> { + var success = failure == null; + var cause = (failure instanceof CompletionException) ? failure.getCause() : failure; + var requestId = (cause instanceof AwsServiceException ase) ? ase.requestId() : null; + + if (!success) { + billingMetrics.recordBatchFailed(batch.size()); + LOGGER.error( + "upload() - error uploading billing to S3, batch:({}), bytes.length:{}, requestId:{}, S3.bucket:{}, S3.key:{}", + batch.description(), + body.length, + requestId, + bucket, + key, + cause); + } else { + billingMetrics.recordBatchDelivered(batch.size()); + LOGGER.debug( + "upload() - success uploading billing to S3, batch:({}), bytes.length:{}, eTag:{}, status:{}, requestId={}, S3.bucket:{}, S3.key:{}", + batch.description(), + body.length, + resp.eTag(), + resp.sdkHttpResponse().statusCode(), + resp.responseMetadata().requestId(), + bucket, + key); + } + return new UploadResult(success, failure, batch); + }); + } + + @Override + public void close() { + client.close(); + } + + private static String objectKey(BatchedLogBuffer.Batch batch) { + + var objectKey = + PATH_PREFIX + + "/" + + OBJECT_KEY_FORMAT.format(batch.oldestEventAt()) + + "/" + + batch.id() + + ".jsonl"; + + if (LOGGER.isTraceEnabled()) { + LOGGER.trace("objectKey() - batch.id:{} , objectKey:{}", batch.id(), objectKey); + } + return objectKey; + } + + private static byte[] objectContent(BatchedLogBuffer.Batch batch) { + + StringBuilder sb = new StringBuilder(); + for (String line : batch.lines()) { + sb.append(line).append('\n'); + } + var bytes = sb.toString().getBytes(StandardCharsets.UTF_8); + + if (LOGGER.isTraceEnabled()) { + LOGGER.trace("objectContent() - batch.id:{} , bytes.length:{}", batch.id(), bytes.length); + } + return bytes; + } +} diff --git a/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingQueueTest.java b/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingQueueTest.java index 98cea178dd..5dc743416d 100644 --- a/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingQueueTest.java +++ b/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingQueueTest.java @@ -17,14 +17,14 @@ import org.junit.jupiter.api.Test; -/** Unit tests for {@link BillingQueue}: seal thresholds, drain limits, and batch metadata. */ +/** Unit tests for {@link BatchedLogBuffer}: seal thresholds, drain limits, and batch metadata. */ class BillingQueueTest { private static final Instant T0 = Instant.parse("2026-05-20T14:23:11Z"); @Test void sealsByEventCount() { - var queue = new BillingQueue(2, 1_000_000, 10); + var queue = new BatchedLogBuffer(2, 1_000_000, 10); queue.offer(T0, "a"); assertThat(queue.shouldFlush()).isFalse(); @@ -35,7 +35,7 @@ void sealsByEventCount() { @Test void sealsByBufferedBytes() { // Each line counts as length + 1 (newline): "aaaa" = 5 bytes. - var queue = new BillingQueue(100, 10, 10); + var queue = new BatchedLogBuffer(100, 10, 10); queue.offer(T0, "aaaa"); assertThat(queue.shouldFlush()).isFalse(); @@ -45,41 +45,41 @@ void sealsByBufferedBytes() { @Test void drainStopsAtMaxEventsAndLeavesTheRemainder() { - var queue = new BillingQueue(2, 1_000_000, 10); + var queue = new BatchedLogBuffer(2, 1_000_000, 10); queue.offer(T0, "a"); queue.offer(T0, "b"); queue.offer(T0, "c"); - assertThat(queue.maybeDrain().lines()).containsExactly("a", "b"); - assertThat(queue.maybeDrain().lines()).containsExactly("c"); - assertThat(queue.maybeDrain().isEmpty()).isTrue(); + assertThat(queue.nextBatch().lines()).containsExactly("a", "b"); + assertThat(queue.nextBatch().lines()).containsExactly("c"); + assertThat(queue.nextBatch().isEmpty()).isTrue(); } @Test void drainStopsAtMaxBytesAndLeavesTheRemainder() { - var queue = new BillingQueue(100, 10, 10); + var queue = new BatchedLogBuffer(100, 10, 10); queue.offer(T0, "aaaa"); queue.offer(T0, "bbbb"); queue.offer(T0, "cccc"); - assertThat(queue.maybeDrain().lines()).containsExactly("aaaa", "bbbb"); - assertThat(queue.maybeDrain().lines()).containsExactly("cccc"); - assertThat(queue.maybeDrain().isEmpty()).isTrue(); + assertThat(queue.nextBatch().lines()).containsExactly("aaaa", "bbbb"); + assertThat(queue.nextBatch().lines()).containsExactly("cccc"); + assertThat(queue.nextBatch().isEmpty()).isTrue(); } @Test void oldestEventAtIsTheMinimumAcrossTheBatchNotTheHead() { - var queue = new BillingQueue(10, 1_000_000, 10); + var queue = new BatchedLogBuffer(10, 1_000_000, 10); // Concurrent publishes can enqueue out of event-time order; the head is not the oldest. queue.offer(T0.plusSeconds(5), "enqueued-first-but-newer"); queue.offer(T0, "enqueued-second-but-older"); - assertThat(queue.maybeDrain().oldestEventAt()).isEqualTo(T0); + assertThat(queue.nextBatch().oldestEventAt()).isEqualTo(T0); } @Test void offerRejectsWhenFull() { - var queue = new BillingQueue(10, 1_000_000, 2); + var queue = new BatchedLogBuffer(10, 1_000_000, 2); assertThat(queue.offer(T0, "a")).isTrue(); assertThat(queue.offer(T0, "b")).isTrue(); @@ -89,7 +89,7 @@ void offerRejectsWhenFull() { @Test void concurrentOfferAndDrainKeepsAccountingConsistent() throws Exception { - var queue = new BillingQueue(10, 1_000_000, 5_000); + var queue = new BatchedLogBuffer(10, 1_000_000, 5_000); int threads = 4; int perThread = 500; Set published = ConcurrentHashMap.newKeySet(); @@ -100,7 +100,7 @@ void concurrentOfferAndDrainKeepsAccountingConsistent() throws Exception { new Thread( () -> { while (!producersDone.get() || !queue.isEmpty()) { - var batch = queue.maybeDrain(); + var batch = queue.nextBatch(); if (batch.isEmpty()) { Thread.onSpinWait(); } else { @@ -150,12 +150,12 @@ void concurrentOfferAndDrainKeepsAccountingConsistent() throws Exception { @Test void byteSealResetsOnceDrained() { - var queue = new BillingQueue(100, 10, 10); + var queue = new BatchedLogBuffer(100, 10, 10); queue.offer(T0, "aaaa"); queue.offer(T0, "bbbb"); assertThat(queue.shouldFlush()).isTrue(); - queue.maybeDrain(); + queue.nextBatch(); assertThat(queue.isEmpty()).isTrue(); assertThat(queue.shouldFlush()).isFalse(); // queuedBytes went back down with the drain diff --git a/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3LogHandlerTest.java b/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3LogHandlerTest.java index 6589edad32..ec8aa6cb97 100644 --- a/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3LogHandlerTest.java +++ b/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3LogHandlerTest.java @@ -82,11 +82,11 @@ static void warmUpMutinyInfrastructure() { // ============================================================ /** - * Programmable {@link BillingS3LogHandler.AsyncBatchUploader}: records every batch and settles + * Programmable {@link AsyncBatchedLogUploader}: records every batch and settles * the returned Uni per {@link Mode}. Never blocks a caller thread — HOLD parks the completion in * {@code held} for the test to release explicitly. */ - static final class RecordingUploader implements BillingS3LogHandler.AsyncBatchUploader { + static final class RecordingUploader implements AsyncBatchedLogUploader { enum Mode { COMPLETE, HOLD, @@ -95,14 +95,14 @@ enum Mode { } volatile Mode mode = Mode.COMPLETE; - final List batches = new CopyOnWriteArrayList<>(); + final List batches = new CopyOnWriteArrayList<>(); final BlockingQueue> held = new LinkedBlockingQueue<>(); final AtomicInteger inFlight = new AtomicInteger(); final AtomicInteger maxInFlight = new AtomicInteger(); volatile boolean closed; @Override - public Uni upload(BillingQueue.Batch batch) { + public Uni upload(BatchedLogBuffer.Batch batch) { batches.add(batch); if (mode == Mode.THROW_SYNC) { throw new RuntimeException("simulated synchronous uploader failure"); diff --git a/src/test/java/io/stargate/sgv2/jsonapi/service/billing/S3BatchUploaderTest.java b/src/test/java/io/stargate/sgv2/jsonapi/service/billing/S3BatchUploaderTest.java index 56d5e70688..9b62fc3f7e 100644 --- a/src/test/java/io/stargate/sgv2/jsonapi/service/billing/S3BatchUploaderTest.java +++ b/src/test/java/io/stargate/sgv2/jsonapi/service/billing/S3BatchUploaderTest.java @@ -36,11 +36,11 @@ class S3BatchUploaderTest { Pattern.compile("data-api/2026/05/20/14/23/[0-9a-f-]{36}\\.jsonl"); private static final String LINE_A = "{\"a\":1}"; private static final String LINE_B = "{\"b\":2}"; - private static final BillingQueue.Batch BATCH = - new BillingQueue.Batch(List.of(LINE_A, LINE_B), Instant.parse("2026-05-20T14:23:11.482Z")); + private static final BatchedLogBuffer.Batch BATCH = + new BatchedLogBuffer.Batch(List.of(LINE_A, LINE_B), Instant.parse("2026-05-20T14:23:11.482Z")); - private static S3BatchUploader uploader(S3AsyncClient client) { - return new S3BatchUploader(client, "my-bucket"); + private static S3BatchedLogUploader uploader(S3AsyncClient client) { + return new S3BatchedLogUploader(client, "my-bucket"); } private static CompletableFuture ok() { @@ -54,14 +54,14 @@ private static CompletableFuture ok() { @Test void objectKeyUsesPathPrefixAndUtcMinutePathFromTimestamp() { var id = UUID.fromString("8c0e9b8a-1d3a-4f6b-9c0d-1234567890ab"); - var key = S3BatchUploader.objectKey(Instant.parse("2026-05-20T14:23:11.482Z"), id); + var key = S3BatchedLogUploader.objectKey(Instant.parse("2026-05-20T14:23:11.482Z"), id); assertThat(key) .isEqualTo("data-api/2026/05/20/14/23/8c0e9b8a-1d3a-4f6b-9c0d-1234567890ab.jsonl"); } @Test void toNdjsonJoinsLinesVerbatimWithTrailingNewlines() { - assertThat(S3BatchUploader.toNdjson(List.of(LINE_A, LINE_B))) + assertThat(S3BatchedLogUploader.toNdjson(List.of(LINE_A, LINE_B))) .isEqualTo((LINE_A + "\n" + LINE_B + "\n").getBytes(StandardCharsets.UTF_8)); } @@ -131,10 +131,10 @@ void closeClosesTheClient() { @Test void createRejectsMissingRegionOrBucket() { - assertThatThrownBy(() -> S3BatchUploader.create(" ", "bucket", Optional.empty())) + assertThatThrownBy(() -> S3BatchedLogUploader.create(" ", "bucket", Optional.empty())) .isInstanceOf(IllegalArgumentException.class) .hasMessageContaining("bucket-region"); - assertThatThrownBy(() -> S3BatchUploader.create("us-east-1", null, Optional.empty())) + assertThatThrownBy(() -> S3BatchedLogUploader.create("us-east-1", null, Optional.empty())) .isInstanceOf(IllegalArgumentException.class) .hasMessageContaining("billing.s3.bucket"); } From 0fcd64f1550a3348cf0d978c5523e7f21c2a13ba Mon Sep 17 00:00:00 2001 From: Aaron Morton Date: Tue, 18 Aug 2026 06:17:00 +1200 Subject: [PATCH 61/65] WIP --- .../metrics/BatchedLogBufferMetrics.java | 70 ++++++++++++++ .../metrics/BatchedLogUploaderMetrics.java | 58 ++++++++++++ .../sgv2/jsonapi/metrics/BillingMetrics.java | 94 ------------------- .../service/billing/BatchedLogBuffer.java | 55 ++++++----- .../billing/BillingS3HandlerInstaller.java | 19 ++-- .../service/billing/BillingS3LogHandler.java | 90 +----------------- .../service/billing/S3BatchedLogUploader.java | 8 +- ...ava => BatchedLogUploaderMetricsTest.java} | 10 +- .../billing/BillingS3LogHandlerTest.java | 4 +- 9 files changed, 175 insertions(+), 233 deletions(-) create mode 100644 src/main/java/io/stargate/sgv2/jsonapi/metrics/BatchedLogBufferMetrics.java create mode 100644 src/main/java/io/stargate/sgv2/jsonapi/metrics/BatchedLogUploaderMetrics.java delete mode 100644 src/main/java/io/stargate/sgv2/jsonapi/metrics/BillingMetrics.java rename src/test/java/io/stargate/sgv2/jsonapi/service/billing/{BillingMetricsTest.java => BatchedLogUploaderMetricsTest.java} (86%) diff --git a/src/main/java/io/stargate/sgv2/jsonapi/metrics/BatchedLogBufferMetrics.java b/src/main/java/io/stargate/sgv2/jsonapi/metrics/BatchedLogBufferMetrics.java new file mode 100644 index 0000000000..a5ce04ddfa --- /dev/null +++ b/src/main/java/io/stargate/sgv2/jsonapi/metrics/BatchedLogBufferMetrics.java @@ -0,0 +1,70 @@ +package io.stargate.sgv2.jsonapi.metrics; + +import io.micrometer.core.instrument.Counter; +import io.micrometer.core.instrument.Gauge; +import io.micrometer.core.instrument.MeterRegistry; +import io.stargate.sgv2.jsonapi.service.billing.BatchedLogBuffer; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.Objects; +import java.util.concurrent.atomic.AtomicBoolean; + +/** + * Metrics for the billing buffer in {@link io.stargate.sgv2.jsonapi.service.billing.BatchedLogBuffer} + */ +public final class BatchedLogBufferMetrics { + + private static final Logger LOGGER = LoggerFactory.getLogger(BatchedLogBufferMetrics.class); + + private final AtomicBoolean bufferRegister = new AtomicBoolean(false); + + private final MeterRegistry meterRegistry; + private final String prefix; + + private final Counter offered; + private final Counter dropped; + + /** + */ + public BatchedLogBufferMetrics(MeterRegistry meterRegistry, String prefix) { + + this.meterRegistry = Objects.requireNonNull(meterRegistry, "meterRegistry must not be null"); + this.prefix = Objects.requireNonNull(prefix, "prefix must not be null"); + if (prefix.isBlank()) { + throw new IllegalArgumentException("prefix must not be blank"); + } + + this.offered = meterRegistry.counter(prefix + ".buffer.offered"); + this.dropped = meterRegistry.counter(prefix + ".buffer.dropped" ); + // Note: not recording events dropped at shutdown as a metric because when shutting down + // the metrics still need to be scrapped to be useful. Do it as a log message that is persistent. + } + + public void registerBuffer(BatchedLogBuffer buffer) { + + if (!bufferRegister.compareAndSet(false, true)) { + throw new IllegalStateException("registerBuffer() already called"); + } + + Gauge.builder(prefix + ".buffer.head_age_ms", + () -> buffer.headEntryAge().toMillis()) + .register(meterRegistry); + Gauge.builder(prefix + ".buffer.size", buffer::size) + .register(meterRegistry); + Gauge.builder(prefix + ".buffer.remaining_capacity", buffer::remainingCapacity) + .register(meterRegistry); + Gauge.builder(prefix + ".buffer.queued_bytes", buffer::queuedBytes) + .register(meterRegistry); + } + + + public void recordOffered() { + offered.increment(); + } + + public void recordDropped() { + dropped.increment(); + } + +} diff --git a/src/main/java/io/stargate/sgv2/jsonapi/metrics/BatchedLogUploaderMetrics.java b/src/main/java/io/stargate/sgv2/jsonapi/metrics/BatchedLogUploaderMetrics.java new file mode 100644 index 0000000000..a576502fd8 --- /dev/null +++ b/src/main/java/io/stargate/sgv2/jsonapi/metrics/BatchedLogUploaderMetrics.java @@ -0,0 +1,58 @@ +package io.stargate.sgv2.jsonapi.metrics; + +import io.micrometer.core.instrument.Counter; +import io.micrometer.core.instrument.Gauge; +import io.micrometer.core.instrument.MeterRegistry; +import java.time.Instant; +import java.util.Objects; +import java.util.concurrent.atomic.AtomicLong; +import java.util.function.Supplier; + +import io.stargate.sgv2.jsonapi.service.billing.BatchedLogBuffer; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Metrics for billing events, mostly around what is sent to S3 + */ +public final class BatchedLogUploaderMetrics { + + private static final Logger LOGGER = LoggerFactory.getLogger(BatchedLogUploaderMetrics.class); + + private final MeterRegistry meterRegistry; + private final String prefix; + + private final Counter batchesUploaded; + private final Counter eventsUploaded; + + private final Counter batchesFailed; + private final Counter eventsFailed; + + /** + */ + public BatchedLogUploaderMetrics(MeterRegistry meterRegistry, String prefix) { + + this.meterRegistry = Objects.requireNonNull(meterRegistry, "meterRegistry must not be null"); + this.prefix = Objects.requireNonNull(prefix, "prefix must not be null"); + if (prefix.isBlank()) { + throw new IllegalArgumentException("prefix must not be blank"); + } + + this.batchesUploaded = meterRegistry.counter(prefix + ".s3.batches.uploaded.size"); + this.batchesUploaded = meterRegistry.counter(prefix + ".s3.batches.uploaded.bytes"); + this.eventsUploaded = meterRegistry.counter(prefix + ".s3.events.uploaded.size"); + + this.batchesFailed = meterRegistry.counter(prefix + ".s3.batches.failed"); + this.eventsFailed = meterRegistry.counter(prefix + ".s3.events.failed"); + } + + public void recordBatchDelivered(BatchedLogBuffer.Batch batch) { + eventsUploaded.increment(size); + batchesUploaded.increment(); + } + + public void recordBatchFailed(BatchedLogBuffer.Batch) { + eventsFailed.increment(size); + batchesFailed.increment(); + } +} diff --git a/src/main/java/io/stargate/sgv2/jsonapi/metrics/BillingMetrics.java b/src/main/java/io/stargate/sgv2/jsonapi/metrics/BillingMetrics.java deleted file mode 100644 index a491bfaf58..0000000000 --- a/src/main/java/io/stargate/sgv2/jsonapi/metrics/BillingMetrics.java +++ /dev/null @@ -1,94 +0,0 @@ -package io.stargate.sgv2.jsonapi.metrics; - -import io.micrometer.core.instrument.Counter; -import io.micrometer.core.instrument.Gauge; -import io.micrometer.core.instrument.MeterRegistry; -import java.time.Instant; -import java.util.concurrent.TimeUnit; -import java.util.concurrent.atomic.AtomicLong; -import java.util.function.Supplier; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; - -/** - * During normal operation, {@code offered = flushed + failed + dropped(capacity) + queue.depth + - * events in in-flight batches}. - * - *

At shutdown, buffered events left after the drain budget are added to {@code - * dropped(shutdown)}. Final counters can be lower than {@code offered} if a concurrent publish - * misses the final queue snapshot or an in-flight upload does not settle before process exit. - * - *

{@code last_delivery.epoch_seconds} is the delivery heartbeat. - */ -public final class BillingMetrics { - - private static final Logger LOG = LoggerFactory.getLogger(BillingMetrics.class); - private static final long DROP_WARN_INTERVAL_NANOS = TimeUnit.MINUTES.toNanos(10); - - // Count of the events sent to the S3 logger - private final Counter offered; - - private final Counter droppedCapacity; - private final Counter droppedShutdown; - private final Counter flushed; - private final Counter failed; - private final Counter batchesUploaded; - private final Counter batchesFailed; - private final AtomicLong lastDeliveryEpochSeconds = new AtomicLong(0); - - /** - * @param depthSource live queue depth, exposed read-only as {@code billing.s3.queue.depth} - */ - public BillingMetrics(MeterRegistry meterRegistry, Supplier depthSource) { - - this.offered = meterRegistry.counter("billing.s3.events.offered"); - this.droppedCapacity = meterRegistry.counter("billing.s3.events.dropped", "reason", "capacity"); - - this.droppedShutdown = meterRegistry.counter("billing.s3.events.dropped", "reason", "shutdown"); - - this.flushed = meterRegistry.counter("billing.s3.events.flushed"); - this.failed = meterRegistry.counter("billing.s3.events.failed"); - this.batchesUploaded = meterRegistry.counter("billing.s3.batches.uploaded"); - this.batchesFailed = meterRegistry.counter("billing.s3.batches.failed"); - - // Catches stalls with no failures to count (e.g. the flush trigger died): alert on staleness - // gated by offered/depth, so idle time isn't mistaken for a dead export. - Gauge.builder( - "billing.s3.last_delivery.epoch_seconds", - lastDeliveryEpochSeconds, - AtomicLong::doubleValue) - .description("Epoch seconds of the last successful batch delivery") - .register(meterRegistry); - Gauge.builder("billing.s3.queue.depth", depthSource) - .description("Billing events buffered in memory, not yet drained") - .register(meterRegistry); - } - - /** A line was handed to the handler (counted before the capacity check). */ - public void recordOffered() { - offered.increment(); - } - - /** A line was dropped on a full buffer; warns rate-limited so a sustained stall stays visible. */ - public void recordDropped() { - droppedCapacity.increment(); - } - - /** Events still buffered when the shutdown budget ran out; close() logs the tombstone. */ - public void recordAbandonedAtShutdown(int size) { - droppedShutdown.increment(size); - } - - /** A batch of event lines landed in S3; bumps the delivery heartbeat. */ - public void recordBatchDelivered(int size) { - flushed.increment(size); - batchesUploaded.increment(); - lastDeliveryEpochSeconds.set(Instant.now().getEpochSecond()); - } - - /** A batch of event lines was given up after the uploader exhausted its retries. */ - public void recordBatchFailed(int size) { - failed.increment(size); - batchesFailed.increment(); - } -} diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BatchedLogBuffer.java b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BatchedLogBuffer.java index cc912fa696..53e1cec52f 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BatchedLogBuffer.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BatchedLogBuffer.java @@ -2,8 +2,8 @@ import com.fasterxml.uuid.Generators; import com.fasterxml.uuid.NoArgGenerator; -import com.google.common.annotations.VisibleForTesting; -import io.stargate.sgv2.jsonapi.metrics.BillingMetrics; +import io.stargate.sgv2.jsonapi.metrics.BatchedLogBufferMetrics; +import io.stargate.sgv2.jsonapi.metrics.BatchedLogUploaderMetrics; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -20,7 +20,7 @@ /** * Buffer for {@link LogRecord} that batches them according to the configuration. *

- * See {@link #BatchedLogBuffer(int, long, Duration, int, BillingMetrics)} for the config. + * See {@link #BatchedLogBuffer(int, long, Duration, int, BatchedLogUploaderMetrics)} for the config. *

*

* There are two uses of this class, producers and consumers. @@ -31,7 +31,7 @@ * * * The buffer is designed to handle these as concurrent calls from different threads, and tracks - * metrics for it's use. + * metrics for its use. *

*/ public class BatchedLogBuffer { @@ -43,7 +43,7 @@ public class BatchedLogBuffer { private final Duration maxAge; private final BlockingQueue queue; - private final BillingMetrics billingMetrics; + private final BatchedLogBufferMetrics metrics; private final AtomicLong queuedBytes = new AtomicLong(0); @@ -57,9 +57,9 @@ public class BatchedLogBuffer { * have many maxBatchBytes if there is a single log record that is bigger. * @param maxAge Maximum age any log record should have in the buffer before a new batch is available. * @param queueCapacity Total number of log records to buffer. - * @param billingMetrics Metrics recording object. + * @param metrics Metrics recording object. */ - public BatchedLogBuffer(int maxBatchSize, long maxBatchBytes, Duration maxAge, int queueCapacity, BillingMetrics billingMetrics) { + public BatchedLogBuffer(int maxBatchSize, long maxBatchBytes, Duration maxAge, int queueCapacity, BatchedLogBufferMetrics metrics) { if (maxBatchSize < 1) { throw new IllegalArgumentException("maxBatchSize must be >= 1, got: " + maxBatchSize ); @@ -77,7 +77,8 @@ public BatchedLogBuffer(int maxBatchSize, long maxBatchBytes, Duration maxAge, i this.maxBytes = maxBatchBytes; this.maxAge = maxAge; - this.billingMetrics = Objects.requireNonNull(billingMetrics, "billingMetrics must not be null"); + this.metrics = Objects.requireNonNull(metrics, "billingMetrics must not be null"); + this.metrics.registerBuffer(this); // must be concurrent to handle multiple threads this.queue = new ArrayBlockingQueue<>(queueCapacity); } @@ -105,10 +106,11 @@ public boolean offer(LogRecord record) { } var newEntry = new Entry(record.getInstant(), logLine); - billingMetrics.recordOffered(); + metrics.recordOffered(); if (!queue.offer(newEntry)) { // Bounded buffer full: drop and count - billingMetrics.recordDropped(); + LOGGER.debug("offer() - buffer full, dropping new entry: {}", newEntry); + metrics.recordDropped(); return false; } @@ -160,7 +162,9 @@ public Batch nextBatch(boolean drainFully) { return null; } - return new Batch(batchReason, lines, oldestEventAt); + LOGGER.info("nextBatch() - next batch of log buffer, reason:{}, lines.size:{}, batchBytes:{}, oldestEventAt: {}", + batchReason, lines.size(), batchBytes, oldestEventAt); + return new Batch(batchReason, lines, batchBytes, oldestEventAt); } public boolean isEmpty() { @@ -171,7 +175,6 @@ public int size() { return queue.size(); } - @VisibleForTesting public long queuedBytes() { return queuedBytes.get(); } @@ -180,29 +183,23 @@ public int remainingCapacity() { return queue.remainingCapacity(); } - - - private Duration oldestEntry() { + public Duration headEntryAge() { var head = queue.peek(); - return head == null ? Duration.ZERO : Duration.between(head.eventAt(), Instant.now()); } - - - private BillingBatchReason decideNextBatch(boolean drainFully) { if (drainFully){ return BillingBatchReason.DRAINING; } if (queue.size() >= maxBatchSize) { - return BillingBatchReason.MAX_BATCH_SIZE_EXCEEDED + return BillingBatchReason.MAX_BATCH_SIZE_EXCEEDED; } if (queuedBytes.get() > maxBytes) { return BillingBatchReason.MAX_BYTES_EXCEEDED; } - if (oldestEntry().compareTo(maxAge) >= 0){ + if (headEntryAge().compareTo(maxAge) >= 0){ return BillingBatchReason.MAX_AGE_EXCEEDED; } return null; @@ -216,24 +213,22 @@ public enum BillingBatchReason { } /** - * OLD BELOW * - * One drained, sealed batch. {@code oldestEventAt} is the minimum event time across {@code lines} - * — queue order is enqueue order, not event-time order, under concurrent publish. */ public static final class Batch { private static final NoArgGenerator UUID_V7_GENERATOR = Generators.timeBasedEpochGenerator(); - private final UUID batchId = UUID_V7_GENERATOR.generate(); private final BillingBatchReason batchReason; private final List lines; + private final long batchBytes; private final Instant oldestEventAt; - public Batch(BillingBatchReason batchReason, List lines, Instant oldestEventAt) { + public Batch(BillingBatchReason batchReason, List lines, long batchBytes, Instant oldestEventAt) { this.batchReason = batchReason; this.lines = Collections.unmodifiableList(lines); + this.batchBytes = batchBytes; this.oldestEventAt = oldestEventAt; } @@ -257,9 +252,13 @@ public int size() { return lines.size(); } + public long batchBytes() { + return batchBytes; + } + public String description() { - return "id:%s, reason:%s, oldestEventAt:%s, size:%s".formatted( - batchId, batchReason, oldestEventAt, size() + return "id:%s, reason:%s, oldestEventAt:%s, size:%s, batchBytes:%s".formatted( + batchId, batchReason, oldestEventAt, size(), batchBytes ); } } diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3HandlerInstaller.java b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3HandlerInstaller.java index 0d4a39109c..c44ff68bf0 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3HandlerInstaller.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3HandlerInstaller.java @@ -11,25 +11,16 @@ import org.slf4j.LoggerFactory; /** - * TODO: XXX MAKE THIS COMMENTS READABLE BY A HUMAN * - *

Attaches a {@link BillingS3LogHandler} to the {@code billing.events} JUL logger at startup + *

Attaches a {@link BillingS3LogHandler} to the {@code billing.events} logger at startup * (when {@link BillingS3ExportConfig#enabled()} is {@code true}) and removes + closes it on * shutdown for a graceful drain. * - *

Done programmatically because Quarkus config can't express it: a category's {@code handlers} - * list can only reference Quarkus's built-in handler types (console/file/syslog/socket), not a - * custom {@link java.util.logging.Handler} class. The one config-driven alternative — a discovered - * {@code @Produces Handler} bean — attaches to the root logger, but {@code billing.events} - * is {@code use-parent-handlers: false} and we want delivery scoped to exactly that category. The - * {@link StartupEvent} observer runs after Quarkus has applied its logging config, so the - * registration sticks. */ @ApplicationScoped public class BillingS3HandlerInstaller { - private static final org.slf4j.Logger LOGGER = - LoggerFactory.getLogger(BillingS3HandlerInstaller.class); + private static final org.slf4j.Logger LOGGER = LoggerFactory.getLogger(BillingS3HandlerInstaller.class); static final String BILLING_LOGGER_NAME = "billing.events"; @@ -47,7 +38,7 @@ public BillingS3HandlerInstaller(BillingS3ExportConfig config, MeterRegistry met void onStart(@Observes StartupEvent event) { if (!config.enabled()) { - LOGGER.debug("Billing S3 export disabled (stargate.jsonapi.billing.s3.enabled=false)"); + LOGGER.info("Billing S3 export disabled (stargate.jsonapi.billing.s3.enabled=false)"); return; } @@ -55,9 +46,11 @@ void onStart(@Observes StartupEvent event) { var bucket = config.bucket().orElse(null); // Fail-loud: invalid billing S3 config throws here, aborting application startup. - var uploader = S3BatchedLogUploader.create(region, bucket, config.endpointOverride()); + var + var uploader = S3BatchedLogUploader.create(region, bucket, config.endpointOverride().orElse(null)); this.handler = new BillingS3LogHandler(config, uploader, meterRegistry); + // TODO: LOGGER NAME SHOULD BE IN CONFIG Logger.getLogger(BILLING_LOGGER_NAME).addHandler(this.handler); diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3LogHandler.java b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3LogHandler.java index ff900c6f26..58c53225eb 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3LogHandler.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3LogHandler.java @@ -2,9 +2,8 @@ import com.google.common.annotations.VisibleForTesting; import io.smallrye.mutiny.Uni; -import io.stargate.sgv2.jsonapi.metrics.BillingMetrics; + import java.time.Duration; -import java.util.Objects; import java.util.concurrent.CountDownLatch; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; @@ -58,16 +57,14 @@ public final class BillingS3LogHandler extends Handler { private final CountDownLatch uploadingFinished = new CountDownLatch(0); private final AsyncBatchedLogUploader uploader; - private final BillingMetrics billingMetrics; private final BatchedLogBuffer batchedLogBuffer; @VisibleForTesting BillingS3LogHandler( - AsyncBatchedLogUploader uploader, BatchedLogBuffer batchedLogBuffer, BillingMetrics billingMetrics) { + AsyncBatchedLogUploader uploader, BatchedLogBuffer batchedLogBuffer) { this.batchedLogBuffer = batchedLogBuffer; this.uploader = uploader; - this.billingMetrics = Objects.requireNonNull(billingMetrics); } private static Duration requirePositive(Duration value, String property) { @@ -178,8 +175,7 @@ void startUploading() { } } } finally { - // record if there are any abandonded events - billingMetrics.recordAbandonedAtShutdown(batchedLogBuffer.size()); + // record if there are any abandoned events if (!batchedLogBuffer.isEmpty()) { LOGGER.warn( "start() - finished with abandoned billing events, billingQueue.size():{} ", @@ -223,84 +219,4 @@ private Uni deferBatch(BatchedLogBuffer.Ba .after(Duration.ofSeconds(10)) .fail(); } - - // /** Seal-triggered flush: ship when the buffer has a full batch by count or bytes. */ - // private void maybeFlush() { - // if (eventQueue.shouldFlush()) { - // tryFlush(); - // } - // } - - // /** - // * Age trigger: every {@code maxAge} tick ships whatever is buffered, sealed or not. - // Deliberately - // * no head-age check: flushing only entries older than {@code maxAge} would let an event that - // just - // * missed a tick wait ~2x{@code maxAge}, while shipping unconditionally bounds every wait by - // one - // * period — at the cost of an occasional small object when a tick lands just after a seal - // flush. - // * - // *

Catches everything: an escaped throwable would silently cancel all future runs of a - // * fixed-rate task. - // */ - // @VisibleForTesting - // void onAgeTick() { - // try { - // if (!eventQueue.isEmpty()) { - // tryFlush(); - // } - // } catch (Throwable t) { - // LOG.error("Billing S3 export age-flush tick failed", t); - // } - // } - // - // /** - // * Claims an in-flight slot (non-blocking CAS, at most {@link #uploadConcurrency} held) and, - // on - // * success, drains + uploads one batch asynchronously. When the upload settles the slot is - // * released and the seal condition re-checked: a full batch may have accumulated meanwhile. - // */ - // private void tryFlush() { - // - // int prev = inFlight.getAndUpdate(n -> n < uploadConcurrency ? n + 1 : n); - // if (prev >= uploadConcurrency) { - // return; - // } - // Uni.createFrom() - // .item(eventQueue::maybeDrain) - // .runSubscriptionOn(Infrastructure.getDefaultWorkerPool()) - // .flatMap(this::uploadBatch) - // .eventually( - // () -> { - // inFlight.getAndDecrement(); - // maybeFlush(); - // }) - // .subscribe() - // .with(ignored -> {}, failure -> LOG.error("Billing S3 export flush failed", failure)); - // } - - /** Uploads one batch; never fails the pipeline — a batch that exhausts retries is counted. */ - // private Uni uploadBatch(BillingQueue.Batch batch) { - // if (batch.isEmpty()) { - // return Uni.createFrom().voidItem(); - // } - // int size = batch.size(); - // // Runs immediately on subscription; deferred only turns a throw before upload() returns a - // Uni - // // into a Uni failure handled below. - // return Uni.createFrom() - // .deferred(() -> uploader.upload(batch)) - // .onItem() - // .invoke(() -> billingMetrics.recordBatchDelivered(size)) - // .onFailure() - // .invoke(t -> LOG.error("Failed to upload billing S3 batch ({} events)", size, t)) - // .onFailure() - // .recoverWithItem( - // () -> { - // billingMetrics.recordBatchFailed(size); - // return null; - // }); - // } - } diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/billing/S3BatchedLogUploader.java b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/S3BatchedLogUploader.java index 72a9bcef7d..07b8bb8e62 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/service/billing/S3BatchedLogUploader.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/S3BatchedLogUploader.java @@ -10,7 +10,7 @@ import java.util.Objects; import java.util.concurrent.CompletionException; -import io.stargate.sgv2.jsonapi.metrics.BillingMetrics; +import io.stargate.sgv2.jsonapi.metrics.BatchedLogUploaderMetrics; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import software.amazon.awssdk.awscore.exception.AwsServiceException; @@ -38,10 +38,10 @@ public class S3BatchedLogUploader implements AsyncBatchedLogUploader { private final S3AsyncClient client; private final String bucket; - private final BillingMetrics billingMetrics; + private final BatchedLogUploaderMetrics billingMetrics; @VisibleForTesting - S3BatchedLogUploader(S3AsyncClient client, String bucket, BillingMetrics billingMetrics) { + S3BatchedLogUploader(S3AsyncClient client, String bucket, BatchedLogUploaderMetrics billingMetrics) { this.client = client; this.bucket = bucket; this.billingMetrics = billingMetrics; @@ -55,7 +55,7 @@ public class S3BatchedLogUploader implements AsyncBatchedLogUploader { * @param endpointOverride * @return */ - public static S3BatchedLogUploader create(String region, String bucket, String endpointOverride, BillingMetrics billingMetrics) { + public static S3BatchedLogUploader create(String region, String bucket, String endpointOverride, BatchedLogUploaderMetrics billingMetrics) { if (region == null || region.isBlank()) { throw new IllegalArgumentException("region must be set"); diff --git a/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingMetricsTest.java b/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BatchedLogUploaderMetricsTest.java similarity index 86% rename from src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingMetricsTest.java rename to src/test/java/io/stargate/sgv2/jsonapi/service/billing/BatchedLogUploaderMetricsTest.java index 47de9e45c6..798e245bd1 100644 --- a/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingMetricsTest.java +++ b/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BatchedLogUploaderMetricsTest.java @@ -6,16 +6,16 @@ import java.time.Instant; import java.util.concurrent.atomic.AtomicInteger; -import io.stargate.sgv2.jsonapi.metrics.BillingMetrics; +import io.stargate.sgv2.jsonapi.metrics.BatchedLogUploaderMetrics; import org.junit.jupiter.api.Test; /** Guards the meter names and tags — dashboards and alerts key on these exact series. */ -class BillingMetricsTest { +class BatchedLogUploaderMetricsTest { @Test void countersFlowToTheExpectedSeries() { var registry = new SimpleMeterRegistry(); - var metrics = new BillingMetrics(registry, () -> 0, 100); + var metrics = new BatchedLogUploaderMetrics(registry, () -> 0, 100); metrics.recordOffered(); metrics.recordDropped(); @@ -38,7 +38,7 @@ void countersFlowToTheExpectedSeries() { void depthGaugeReadsTheLiveSupplier() { var registry = new SimpleMeterRegistry(); var depth = new AtomicInteger(7); - new BillingMetrics(registry, depth::get, 100); + new BatchedLogUploaderMetrics(registry, depth::get, 100); assertThat(registry.get("billing.s3.queue.depth").gauge().value()).isEqualTo(7.0); depth.set(11); @@ -48,7 +48,7 @@ void depthGaugeReadsTheLiveSupplier() { @Test void deliveryHeartbeatAdvancesOnDeliveredBatches() { var registry = new SimpleMeterRegistry(); - var metrics = new BillingMetrics(registry, () -> 0, 100); + var metrics = new BatchedLogUploaderMetrics(registry, () -> 0, 100); var heartbeat = registry.get("billing.s3.last_delivery.epoch_seconds").gauge(); assertThat(heartbeat.value()).isZero(); // never delivered diff --git a/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3LogHandlerTest.java b/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3LogHandlerTest.java index ec8aa6cb97..ff5ddab5f1 100644 --- a/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3LogHandlerTest.java +++ b/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3LogHandlerTest.java @@ -27,7 +27,7 @@ import java.util.logging.Level; import java.util.logging.LogRecord; -import io.stargate.sgv2.jsonapi.metrics.BillingMetrics; +import io.stargate.sgv2.jsonapi.metrics.BatchedLogUploaderMetrics; import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.Test; import org.slf4j.Logger; @@ -634,7 +634,7 @@ void overflowAccountingReconciles() throws Exception { * returns instead of hanging, delivered lines are a duplicate-free subset of what was published, * and the metrics reconcile as {@code flushed + dropped <= offered} rather than {@code ==}. The * gap is expected, not a bug: a publish can land after close() takes its final buffer snapshot, - * so that line is neither delivered nor counted as dropped — see {@link BillingMetrics}'s class + * so that line is neither delivered nor counted as dropped — see {@link BatchedLogUploaderMetrics}'s class * doc for this same at-most-once slippage. The log line below reports the exact gap each run. */ @Test From b1a3d2f24d9adc791a520270953cd56747d47ec8 Mon Sep 17 00:00:00 2001 From: Aaron Morton Date: Thu, 20 Aug 2026 06:46:54 +1200 Subject: [PATCH 62/65] WIP --- .../jsonapi/config/BillingS3ExportConfig.java | 21 +--- .../metrics/BatchedLogBufferMetrics.java | 45 ++----- .../metrics/BatchedLogUploaderMetrics.java | 60 +++++---- .../sgv2/jsonapi/metrics/MetricsBase.java | 67 ++++++++++ .../service/billing/BatchedLogBuffer.java | 8 +- .../billing/BillingS3HandlerInstaller.java | 24 ++-- .../service/billing/S3BatchedLogUploader.java | 114 +++++++++--------- .../BillingS3HandlerInstallerTest.java | 2 +- 8 files changed, 192 insertions(+), 149 deletions(-) create mode 100644 src/main/java/io/stargate/sgv2/jsonapi/metrics/MetricsBase.java diff --git a/src/main/java/io/stargate/sgv2/jsonapi/config/BillingS3ExportConfig.java b/src/main/java/io/stargate/sgv2/jsonapi/config/BillingS3ExportConfig.java index 726f4ac584..3aade93cec 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/config/BillingS3ExportConfig.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/config/BillingS3ExportConfig.java @@ -13,11 +13,11 @@ public interface BillingS3ExportConfig { @WithDefault("false") boolean enabled(); - /** S3 bucket name */ - Optional bucket(); - /** S3 bucket region */ - Optional bucketRegion(); + String region(); + + /** S3 bucket name */ + String bucket(); /** * Only for non-AWS S3 endpoints (e.g. S3Mock in tests). TODO: XXX EXPLAIN WHAT THIS SHOULD SET @@ -26,7 +26,7 @@ public interface BillingS3ExportConfig { Optional endpointOverride(); /** */ - @WithDefault("50") + @WithDefault("2048") int maxEventsPerBatch(); /** @@ -43,15 +43,4 @@ public interface BillingS3ExportConfig { /** Bound on buffered events; beyond it new lines are dropped. */ @WithDefault("10000") int queueCapacity(); - - /** Max concurrent S3 PUTs. */ - @WithDefault("4") - int uploadConcurrency(); - - /** - * Budget for draining the buffer at shutdown; keep below the pod termination grace period. TODO: - * XXX WHAT IS THE CURRENT TERMINATION PERIOD ? - */ - @WithDefault("PT20S") - Duration shutdownTimeout(); } diff --git a/src/main/java/io/stargate/sgv2/jsonapi/metrics/BatchedLogBufferMetrics.java b/src/main/java/io/stargate/sgv2/jsonapi/metrics/BatchedLogBufferMetrics.java index a5ce04ddfa..9c6c720d15 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/metrics/BatchedLogBufferMetrics.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/metrics/BatchedLogBufferMetrics.java @@ -8,35 +8,27 @@ import org.slf4j.LoggerFactory; import java.util.Objects; +import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; /** * Metrics for the billing buffer in {@link io.stargate.sgv2.jsonapi.service.billing.BatchedLogBuffer} */ -public final class BatchedLogBufferMetrics { - - private static final Logger LOGGER = LoggerFactory.getLogger(BatchedLogBufferMetrics.class); +public final class BatchedLogBufferMetrics extends MetricsBase { private final AtomicBoolean bufferRegister = new AtomicBoolean(false); - private final MeterRegistry meterRegistry; - private final String prefix; - - private final Counter offered; - private final Counter dropped; + public final Counter offered; + public final Counter dropped; /** */ public BatchedLogBufferMetrics(MeterRegistry meterRegistry, String prefix) { + super(meterRegistry, prefix); - this.meterRegistry = Objects.requireNonNull(meterRegistry, "meterRegistry must not be null"); - this.prefix = Objects.requireNonNull(prefix, "prefix must not be null"); - if (prefix.isBlank()) { - throw new IllegalArgumentException("prefix must not be blank"); - } + this.offered = newCounter("buffer.offered"); + this.dropped = newCounter("buffer.dropped"); - this.offered = meterRegistry.counter(prefix + ".buffer.offered"); - this.dropped = meterRegistry.counter(prefix + ".buffer.dropped" ); // Note: not recording events dropped at shutdown as a metric because when shutting down // the metrics still need to be scrapped to be useful. Do it as a log message that is persistent. } @@ -47,24 +39,9 @@ public void registerBuffer(BatchedLogBuffer buffer) { throw new IllegalStateException("registerBuffer() already called"); } - Gauge.builder(prefix + ".buffer.head_age_ms", - () -> buffer.headEntryAge().toMillis()) - .register(meterRegistry); - Gauge.builder(prefix + ".buffer.size", buffer::size) - .register(meterRegistry); - Gauge.builder(prefix + ".buffer.remaining_capacity", buffer::remainingCapacity) - .register(meterRegistry); - Gauge.builder(prefix + ".buffer.queued_bytes", buffer::queuedBytes) - .register(meterRegistry); + newTimeGauge("buffer.head_age_ms",() -> buffer.headEntryAge().toMillis() , TimeUnit.MILLISECONDS); + newGauge("buffer.size", buffer::size); + newGauge("buffer.remaining_capacity", buffer::remainingCapacity); + newGauge("buffer.bytes", buffer::queuedBytes); } - - - public void recordOffered() { - offered.increment(); - } - - public void recordDropped() { - dropped.increment(); - } - } diff --git a/src/main/java/io/stargate/sgv2/jsonapi/metrics/BatchedLogUploaderMetrics.java b/src/main/java/io/stargate/sgv2/jsonapi/metrics/BatchedLogUploaderMetrics.java index a576502fd8..9552fb8b6f 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/metrics/BatchedLogUploaderMetrics.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/metrics/BatchedLogUploaderMetrics.java @@ -1,58 +1,54 @@ package io.stargate.sgv2.jsonapi.metrics; import io.micrometer.core.instrument.Counter; -import io.micrometer.core.instrument.Gauge; import io.micrometer.core.instrument.MeterRegistry; -import java.time.Instant; -import java.util.Objects; -import java.util.concurrent.atomic.AtomicLong; -import java.util.function.Supplier; +import io.micrometer.core.instrument.Timer; import io.stargate.sgv2.jsonapi.service.billing.BatchedLogBuffer; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; + +import java.util.Objects; /** * Metrics for billing events, mostly around what is sent to S3 */ -public final class BatchedLogUploaderMetrics { - - private static final Logger LOGGER = LoggerFactory.getLogger(BatchedLogUploaderMetrics.class); +public final class BatchedLogUploaderMetrics extends MetricsBase{ - private final MeterRegistry meterRegistry; - private final String prefix; - private final Counter batchesUploaded; - private final Counter eventsUploaded; + public final Counter uploadedBatches; + public final Counter uploadedBytes; + public final Timer uploadedHeadAgeMs; + public final Counter uploadedEvents; - private final Counter batchesFailed; - private final Counter eventsFailed; + public final Counter failedBatches; + public final Counter failedEvents; /** */ public BatchedLogUploaderMetrics(MeterRegistry meterRegistry, String prefix) { + super(meterRegistry, prefix); - this.meterRegistry = Objects.requireNonNull(meterRegistry, "meterRegistry must not be null"); - this.prefix = Objects.requireNonNull(prefix, "prefix must not be null"); - if (prefix.isBlank()) { - throw new IllegalArgumentException("prefix must not be blank"); - } + this.uploadedBatches = newCounter("s3.uploaded.batches"); + this.uploadedBytes = newCounter("s3.uploaded.bytes"); + this.uploadedEvents = newCounter("s3.uploaded.events"); + this.uploadedHeadAgeMs = newTimer("s3.uploaded.oldest_event"); - this.batchesUploaded = meterRegistry.counter(prefix + ".s3.batches.uploaded.size"); - this.batchesUploaded = meterRegistry.counter(prefix + ".s3.batches.uploaded.bytes"); - this.eventsUploaded = meterRegistry.counter(prefix + ".s3.events.uploaded.size"); - - this.batchesFailed = meterRegistry.counter(prefix + ".s3.batches.failed"); - this.eventsFailed = meterRegistry.counter(prefix + ".s3.events.failed"); + this.failedBatches = newCounter( "s3.failed.batches"); + this.failedEvents = newCounter( "s3.failed.events"); } public void recordBatchDelivered(BatchedLogBuffer.Batch batch) { - eventsUploaded.increment(size); - batchesUploaded.increment(); + + Objects.requireNonNull(batch, "batch must not be null"); + uploadedBatches.increment(); + uploadedBytes.increment(batch.batchBytes()); + uploadedHeadAgeMs.record(batch.oldestEventAtDuration()); + uploadedEvents.increment(batch.size()); } - public void recordBatchFailed(BatchedLogBuffer.Batch) { - eventsFailed.increment(size); - batchesFailed.increment(); + public void recordBatchFailed(BatchedLogBuffer.Batch batch) { + + Objects.requireNonNull(batch, "batch must not be null"); + failedBatches.increment(); + failedEvents.increment(batch.size()); } } diff --git a/src/main/java/io/stargate/sgv2/jsonapi/metrics/MetricsBase.java b/src/main/java/io/stargate/sgv2/jsonapi/metrics/MetricsBase.java new file mode 100644 index 0000000000..6a00cb1f88 --- /dev/null +++ b/src/main/java/io/stargate/sgv2/jsonapi/metrics/MetricsBase.java @@ -0,0 +1,67 @@ +package io.stargate.sgv2.jsonapi.metrics; + +import io.micrometer.core.instrument.*; + +import java.util.Objects; +import java.util.concurrent.TimeUnit; +import java.util.function.Supplier; + +/** + * Common base for classes that create metric measures + */ +public abstract class MetricsBase { + + protected final MeterRegistry meterRegistry; + protected final String prefix; + + protected MetricsBase(MeterRegistry meterRegistry, String prefix) { + + this.meterRegistry = Objects.requireNonNull(meterRegistry, "meterRegistry must not be null"); + this.prefix = Objects.requireNonNull(prefix, "prefix must not be null"); + if (prefix.isBlank()) { + throw new IllegalArgumentException("prefix must not be blank"); + } + } + + protected String validateName(String name){ + Objects.requireNonNull(name, "name must not be null"); + if(name.isBlank()){ + throw new IllegalArgumentException("name must not be blank"); + } + + return name.charAt(0) == '.' ? name : "." + name; + } + + protected String fullName(String name){ + return prefix + validateName(name); + } + + protected Counter newCounter(String name) { + return meterRegistry.counter(fullName(name)); + } + + protected Gauge newGauge(String name, Supplier func) { + // no null checks in the builder below + Objects.requireNonNull(func, "func must not be null"); + + return Gauge.builder(fullName(name), func) + .register(meterRegistry); + } + + protected Timer newTimer(String name){ + return newTimer(name, 0.5, 0.95, 0.99); + } + + protected Timer newTimer(String name, double... percentiles) { + return Timer.builder(fullName(name)) + .publishPercentiles(percentiles) + .register(meterRegistry); + } + + + protected TimeGauge newTimeGauge(String name, Supplier func, TimeUnit unit) { + + return TimeGauge.builder(fullName(name), func, unit) + .register(meterRegistry ); + } +} diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BatchedLogBuffer.java b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BatchedLogBuffer.java index 53e1cec52f..ee1efa3a58 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BatchedLogBuffer.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BatchedLogBuffer.java @@ -106,11 +106,11 @@ public boolean offer(LogRecord record) { } var newEntry = new Entry(record.getInstant(), logLine); - metrics.recordOffered(); + metrics.offered.increment(); if (!queue.offer(newEntry)) { // Bounded buffer full: drop and count LOGGER.debug("offer() - buffer full, dropping new entry: {}", newEntry); - metrics.recordDropped(); + metrics.dropped.increment(); return false; } @@ -248,6 +248,10 @@ public Instant oldestEventAt() { return oldestEventAt; } + public Duration oldestEventAtDuration() { + return Duration.between(oldestEventAt(), Instant.now()); + } + public int size() { return lines.size(); } diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3HandlerInstaller.java b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3HandlerInstaller.java index c44ff68bf0..384fc8796f 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3HandlerInstaller.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3HandlerInstaller.java @@ -4,6 +4,8 @@ import io.quarkus.runtime.ShutdownEvent; import io.quarkus.runtime.StartupEvent; import io.stargate.sgv2.jsonapi.config.BillingS3ExportConfig; +import io.stargate.sgv2.jsonapi.metrics.BatchedLogBufferMetrics; +import io.stargate.sgv2.jsonapi.metrics.BatchedLogUploaderMetrics; import jakarta.enterprise.context.ApplicationScoped; import jakarta.enterprise.event.Observes; import jakarta.inject.Inject; @@ -22,7 +24,8 @@ public class BillingS3HandlerInstaller { private static final org.slf4j.Logger LOGGER = LoggerFactory.getLogger(BillingS3HandlerInstaller.class); - static final String BILLING_LOGGER_NAME = "billing.events"; + private static final String METRICS_PREFIX ="billing"; + private static final String BILLING_LOGGER_NAME = "billing.events"; private final BillingS3ExportConfig config; private final MeterRegistry meterRegistry; @@ -38,16 +41,23 @@ public BillingS3HandlerInstaller(BillingS3ExportConfig config, MeterRegistry met void onStart(@Observes StartupEvent event) { if (!config.enabled()) { - LOGGER.info("Billing S3 export disabled (stargate.jsonapi.billing.s3.enabled=false)"); + LOGGER.info("Billing S3 export disabled"); return; } - var region = config.bucketRegion().orElse(null); - var bucket = config.bucket().orElse(null); - // Fail-loud: invalid billing S3 config throws here, aborting application startup. - var - var uploader = S3BatchedLogUploader.create(region, bucket, config.endpointOverride().orElse(null)); + var uploader = S3BatchedLogUploader.create( + config.region(), + config.bucket(), + config.endpointOverride().orElse(null), + new BatchedLogUploaderMetrics(meterRegistry, METRICS_PREFIX)); + + var buffer= new BatchedLogBuffer( + config.maxEventsPerBatch(), + config.maxBytesPerBatch(), + config.maxAge(), + config.queueCapacity(), + new BatchedLogBufferMetrics(meterRegistry, METRICS_PREFIX)); this.handler = new BillingS3LogHandler(config, uploader, meterRegistry); diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/billing/S3BatchedLogUploader.java b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/S3BatchedLogUploader.java index 07b8bb8e62..5eec6d1fd7 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/service/billing/S3BatchedLogUploader.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/S3BatchedLogUploader.java @@ -19,7 +19,11 @@ import software.amazon.awssdk.services.s3.S3AsyncClient; import software.amazon.awssdk.services.s3.model.PutObjectRequest; -/** Uploads sealed billing batches to S3 as NDJSON objects under time-partitioned keys. */ +/** Uploads sealed billing batches to S3 as NDJSON objects under time-partitioned keys. + * TODO: + * .requestChecksumCalculation(RequestChecksumCalculation.WHEN_SUPPORTED) + * .responseChecksumValidation(ResponseChecksumValidation.WHEN_SUPPORTED) + * */ public class S3BatchedLogUploader implements AsyncBatchedLogUploader { private static final Logger LOGGER = LoggerFactory.getLogger(S3BatchedLogUploader.class); @@ -27,24 +31,27 @@ public class S3BatchedLogUploader implements AsyncBatchedLogUploader { // S3 destination formatting private static final String PATH_PREFIX = "data-api"; private static final String CONTENT_TYPE_NDJSON = "application/x-ndjson"; - private static final DateTimeFormatter OBJECT_KEY_FORMAT = + private static final DateTimeFormatter OBJECT_KEY_FORMATTER = DateTimeFormatter.ofPattern("yyyy/MM/dd/HH/mm").withZone(ZoneOffset.UTC); - // Bound every PUT so a hung connection can neither pin an upload slot indefinitely nor stall - // the shutdown drain. Retries stay inside the SDK's built-in default policy (bounded attempts, - // jittered throttle-aware backoff). + private static final Duration API_CALL_ATTEMPT_TIMEOUT = Duration.ofSeconds(10); private static final Duration API_CALL_TIMEOUT = Duration.ofSeconds(30); private final S3AsyncClient client; + private final String region; private final String bucket; - private final BatchedLogUploaderMetrics billingMetrics; + private final BatchedLogUploaderMetrics metrics; + - @VisibleForTesting - S3BatchedLogUploader(S3AsyncClient client, String bucket, BatchedLogUploaderMetrics billingMetrics) { + private S3BatchedLogUploader(S3AsyncClient client, + String region, + String bucket, + BatchedLogUploaderMetrics metrics) { this.client = client; + this.region = region; this.bucket = bucket; - this.billingMetrics = billingMetrics; + this.metrics = metrics; } /** @@ -55,22 +62,25 @@ public class S3BatchedLogUploader implements AsyncBatchedLogUploader { * @param endpointOverride * @return */ - public static S3BatchedLogUploader create(String region, String bucket, String endpointOverride, BatchedLogUploaderMetrics billingMetrics) { + public static S3BatchedLogUploader create(String region, + String bucket, + String endpointOverride, + BatchedLogUploaderMetrics metrics) { if (region == null || region.isBlank()) { - throw new IllegalArgumentException("region must be set"); + throw new IllegalArgumentException("region must not be null or blank"); } if (bucket == null || bucket.isBlank()) { - throw new IllegalArgumentException("bucket must be set"); + throw new IllegalArgumentException("bucket must not be null or blank"); } + Objects.requireNonNull(metrics, "metrics must not be null"); // Credentials resolve from the SDK's default provider chain (env vars, web-identity/OIDC // token, instance/container roles), left implicit so the client owns — and closes — the // provider. This transparently supports federated (AssumeRoleWithWebIdentity) and // cross-account access: the bucket may live in a different account (per IAM + bucket // policy); its region is set via .region(). - var builder = - S3AsyncClient.builder() + var builder = S3AsyncClient.builder() .region(Region.of(region)) .overrideConfiguration( config -> @@ -86,7 +96,7 @@ public static S3BatchedLogUploader create(String region, String bucket, String e builder.endpointOverride(URI.create(endpointOverride)).forcePathStyle(true); } - return new S3BatchedLogUploader(builder.build(), bucket); + return new S3BatchedLogUploader(builder.build(), region, bucket, metrics); } @Override @@ -94,32 +104,29 @@ public Uni upload(BatchedLogBuffer.Batch batch) { Objects.requireNonNull(batch, "batch must not be null"); - var key = objectKey(batch); + var location = objectLocation(batch); var body = objectContent(batch); - if (LOGGER.isDebugEnabled()) { - LOGGER.debug( - "upload() - starting to upload batch, batch:({}), S3.bucket:{}, S3.key:{}", - batch.description(), - bucket, - key); - } + + LOGGER.info( + "upload() - starting to upload batch, batch:({}), location:{}, body.size:{}", + batch.description(), + location, + body.length); // No .retry() here: unconfigured, S3AsyncClient already retries (default LegacyRetryStrategy — // 3 retries / 4 attempts). See // https://docs.aws.amazon.com/sdk-for-java/latest/developer-guide/retry-strategy.html // and see https://github.com/aws/aws-sdk-java-v2/issues/6987 for future change. + var putRequest = PutObjectRequest.builder() + .bucket(location.bucket()) + .key(location.key()) + .contentType(CONTENT_TYPE_NDJSON) + .build(); + return Uni.createFrom() - .completionStage( - () -> - client.putObject( - PutObjectRequest.builder() - .bucket(bucket) - .key(key) - .contentType(CONTENT_TYPE_NDJSON) - .build(), - AsyncRequestBody.fromBytes(body))) + .completionStage(client.putObject(putRequest, AsyncRequestBody.fromBytes(body))) .onItemOrFailure() .transform( (resp, failure) -> { @@ -128,26 +135,22 @@ public Uni upload(BatchedLogBuffer.Batch batch) { var requestId = (cause instanceof AwsServiceException ase) ? ase.requestId() : null; if (!success) { - billingMetrics.recordBatchFailed(batch.size()); + metrics.recordBatchFailed(batch); LOGGER.error( - "upload() - error uploading billing to S3, batch:({}), bytes.length:{}, requestId:{}, S3.bucket:{}, S3.key:{}", + "upload() - error uploading billing to S3, batch:({}), location:{}, requestId:{}", batch.description(), - body.length, + location, requestId, - bucket, - key, cause); } else { - billingMetrics.recordBatchDelivered(batch.size()); - LOGGER.debug( - "upload() - success uploading billing to S3, batch:({}), bytes.length:{}, eTag:{}, status:{}, requestId={}, S3.bucket:{}, S3.key:{}", + metrics.recordBatchDelivered(batch); + LOGGER.info( + "upload() - success uploading billing to S3, batch:({}), location:{}, requestId:{}, eTag:{}, status:{}", batch.description(), - body.length, + location, + requestId, resp.eTag(), - resp.sdkHttpResponse().statusCode(), - resp.responseMetadata().requestId(), - bucket, - key); + resp.sdkHttpResponse().statusCode()); } return new UploadResult(success, failure, batch); }); @@ -158,20 +161,17 @@ public void close() { client.close(); } - private static String objectKey(BatchedLogBuffer.Batch batch) { + private S3Location objectLocation(BatchedLogBuffer.Batch batch) { var objectKey = PATH_PREFIX + "/" - + OBJECT_KEY_FORMAT.format(batch.oldestEventAt()) + + OBJECT_KEY_FORMATTER.format(batch.oldestEventAt()) + "/" + batch.id() + ".jsonl"; - if (LOGGER.isTraceEnabled()) { - LOGGER.trace("objectKey() - batch.id:{} , objectKey:{}", batch.id(), objectKey); - } - return objectKey; + return new S3Location(region, bucket, objectKey); } private static byte[] objectContent(BatchedLogBuffer.Batch batch) { @@ -180,11 +180,11 @@ private static byte[] objectContent(BatchedLogBuffer.Batch batch) { for (String line : batch.lines()) { sb.append(line).append('\n'); } - var bytes = sb.toString().getBytes(StandardCharsets.UTF_8); - - if (LOGGER.isTraceEnabled()) { - LOGGER.trace("objectContent() - batch.id:{} , bytes.length:{}", batch.id(), bytes.length); - } - return bytes; + return sb.toString().getBytes(StandardCharsets.UTF_8); } + + private record S3Location( + String region, + String bucket, + String key) {} } diff --git a/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3HandlerInstallerTest.java b/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3HandlerInstallerTest.java index 33f85ebf1b..1e4f9405db 100644 --- a/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3HandlerInstallerTest.java +++ b/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3HandlerInstallerTest.java @@ -27,7 +27,7 @@ private static BillingS3ExportConfig config(boolean enabled, String bucket, Stri BillingS3ExportConfig config = mock(BillingS3ExportConfig.class); when(config.enabled()).thenReturn(enabled); when(config.bucket()).thenReturn(Optional.ofNullable(bucket)); - when(config.bucketRegion()).thenReturn(Optional.ofNullable(region)); + when(config.region()).thenReturn(Optional.ofNullable(region)); when(config.endpointOverride()).thenReturn(Optional.empty()); when(config.maxEventsPerBatch()).thenReturn(50); when(config.maxBytesPerBatch()).thenReturn(2_097_152L); From b167eb4377dabfdbaf8d115e2b3f34528cab8dc7 Mon Sep 17 00:00:00 2001 From: Aaron Morton Date: Tue, 1 Sep 2026 17:10:46 +1200 Subject: [PATCH 63/65] WIP - BatchedLogBuffer & BatchedLogBufferTest done --- .../jsonapi/config/BillingS3ExportConfig.java | 2 + .../metrics/BatchedLogBufferMetrics.java | 29 +- .../metrics/BatchedLogUploaderMetrics.java | 18 +- .../sgv2/jsonapi/metrics/MetricsBase.java | 78 +- .../billing/AsyncBatchedLogUploader.java | 41 +- .../service/billing/BatchedLogBuffer.java | 366 ++-- .../billing/BillingS3HandlerInstaller.java | 34 +- .../service/billing/BillingS3LogHandler.java | 11 +- .../service/billing/S3BatchedLogUploader.java | 74 +- .../stargate/sgv2/jsonapi/TestConstants.java | 2 +- .../service/billing/BatchedLogBufferTest.java | 913 ++++++++++ .../BatchedLogUploaderMetricsTest.java | 101 +- .../service/billing/BillingEventTest.java | 1 - .../service/billing/BillingQueueTest.java | 163 -- .../BillingS3HandlerInstallerTest.java | 135 +- .../billing/BillingS3LogHandlerTest.java | 1465 ++++++++--------- .../service/billing/DefaultBillingTest.java | 9 +- .../service/billing/S3BatchUploaderTest.java | 245 ++- .../stargate/sgv2/jsonapi/util/MockClock.java | 62 + 19 files changed, 2320 insertions(+), 1429 deletions(-) create mode 100644 src/test/java/io/stargate/sgv2/jsonapi/service/billing/BatchedLogBufferTest.java delete mode 100644 src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingQueueTest.java create mode 100644 src/test/java/io/stargate/sgv2/jsonapi/util/MockClock.java diff --git a/src/main/java/io/stargate/sgv2/jsonapi/config/BillingS3ExportConfig.java b/src/main/java/io/stargate/sgv2/jsonapi/config/BillingS3ExportConfig.java index 3aade93cec..ef2d7ba3a4 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/config/BillingS3ExportConfig.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/config/BillingS3ExportConfig.java @@ -14,9 +14,11 @@ public interface BillingS3ExportConfig { boolean enabled(); /** S3 bucket region */ + @WithDefault("us-east-2") String region(); /** S3 bucket name */ + @WithDefault("serverless-usage-dev") String bucket(); /** diff --git a/src/main/java/io/stargate/sgv2/jsonapi/metrics/BatchedLogBufferMetrics.java b/src/main/java/io/stargate/sgv2/jsonapi/metrics/BatchedLogBufferMetrics.java index 9c6c720d15..a38bef1f9d 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/metrics/BatchedLogBufferMetrics.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/metrics/BatchedLogBufferMetrics.java @@ -1,28 +1,23 @@ package io.stargate.sgv2.jsonapi.metrics; import io.micrometer.core.instrument.Counter; -import io.micrometer.core.instrument.Gauge; import io.micrometer.core.instrument.MeterRegistry; import io.stargate.sgv2.jsonapi.service.billing.BatchedLogBuffer; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; - -import java.util.Objects; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; /** - * Metrics for the billing buffer in {@link io.stargate.sgv2.jsonapi.service.billing.BatchedLogBuffer} + * Metrics for the billing buffer in {@link + * io.stargate.sgv2.jsonapi.service.billing.BatchedLogBuffer} */ public final class BatchedLogBufferMetrics extends MetricsBase { private final AtomicBoolean bufferRegister = new AtomicBoolean(false); - public final Counter offered; - public final Counter dropped; + private final Counter offered; + private final Counter dropped; - /** - */ + /** */ public BatchedLogBufferMetrics(MeterRegistry meterRegistry, String prefix) { super(meterRegistry, prefix); @@ -30,7 +25,8 @@ public BatchedLogBufferMetrics(MeterRegistry meterRegistry, String prefix) { this.dropped = newCounter("buffer.dropped"); // Note: not recording events dropped at shutdown as a metric because when shutting down - // the metrics still need to be scrapped to be useful. Do it as a log message that is persistent. + // the metrics still need to be scrapped to be useful. Do it as a log message that is + // persistent. } public void registerBuffer(BatchedLogBuffer buffer) { @@ -39,9 +35,18 @@ public void registerBuffer(BatchedLogBuffer buffer) { throw new IllegalStateException("registerBuffer() already called"); } - newTimeGauge("buffer.head_age_ms",() -> buffer.headEntryAge().toMillis() , TimeUnit.MILLISECONDS); + newTimeGauge( + "buffer.head_age_ms", () -> buffer.headEntryAge().toMillis(), TimeUnit.MILLISECONDS); newGauge("buffer.size", buffer::size); newGauge("buffer.remaining_capacity", buffer::remainingCapacity); newGauge("buffer.bytes", buffer::queuedBytes); } + + public void offered() { + offered.increment(); + } + + public void dropped() { + dropped.increment(); + } } diff --git a/src/main/java/io/stargate/sgv2/jsonapi/metrics/BatchedLogUploaderMetrics.java b/src/main/java/io/stargate/sgv2/jsonapi/metrics/BatchedLogUploaderMetrics.java index 9552fb8b6f..608970af5a 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/metrics/BatchedLogUploaderMetrics.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/metrics/BatchedLogUploaderMetrics.java @@ -2,17 +2,12 @@ import io.micrometer.core.instrument.Counter; import io.micrometer.core.instrument.MeterRegistry; - import io.micrometer.core.instrument.Timer; import io.stargate.sgv2.jsonapi.service.billing.BatchedLogBuffer; - import java.util.Objects; -/** - * Metrics for billing events, mostly around what is sent to S3 - */ -public final class BatchedLogUploaderMetrics extends MetricsBase{ - +/** Metrics for billing events, mostly around what is sent to S3 */ +public final class BatchedLogUploaderMetrics extends MetricsBase { public final Counter uploadedBatches; public final Counter uploadedBytes; @@ -22,8 +17,7 @@ public final class BatchedLogUploaderMetrics extends MetricsBase{ public final Counter failedBatches; public final Counter failedEvents; - /** - */ + /** */ public BatchedLogUploaderMetrics(MeterRegistry meterRegistry, String prefix) { super(meterRegistry, prefix); @@ -32,15 +26,15 @@ public BatchedLogUploaderMetrics(MeterRegistry meterRegistry, String prefix) { this.uploadedEvents = newCounter("s3.uploaded.events"); this.uploadedHeadAgeMs = newTimer("s3.uploaded.oldest_event"); - this.failedBatches = newCounter( "s3.failed.batches"); - this.failedEvents = newCounter( "s3.failed.events"); + this.failedBatches = newCounter("s3.failed.batches"); + this.failedEvents = newCounter("s3.failed.events"); } public void recordBatchDelivered(BatchedLogBuffer.Batch batch) { Objects.requireNonNull(batch, "batch must not be null"); uploadedBatches.increment(); - uploadedBytes.increment(batch.batchBytes()); + uploadedBytes.increment(batch.bytes()); uploadedHeadAgeMs.record(batch.oldestEventAtDuration()); uploadedEvents.increment(batch.size()); } diff --git a/src/main/java/io/stargate/sgv2/jsonapi/metrics/MetricsBase.java b/src/main/java/io/stargate/sgv2/jsonapi/metrics/MetricsBase.java index 6a00cb1f88..76461da3e8 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/metrics/MetricsBase.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/metrics/MetricsBase.java @@ -1,67 +1,59 @@ package io.stargate.sgv2.jsonapi.metrics; import io.micrometer.core.instrument.*; - import java.util.Objects; import java.util.concurrent.TimeUnit; import java.util.function.Supplier; -/** - * Common base for classes that create metric measures - */ +/** Common base for classes that create metric measures */ public abstract class MetricsBase { - protected final MeterRegistry meterRegistry; - protected final String prefix; + protected final MeterRegistry meterRegistry; + protected final String prefix; - protected MetricsBase(MeterRegistry meterRegistry, String prefix) { + protected MetricsBase(MeterRegistry meterRegistry, String prefix) { - this.meterRegistry = Objects.requireNonNull(meterRegistry, "meterRegistry must not be null"); - this.prefix = Objects.requireNonNull(prefix, "prefix must not be null"); - if (prefix.isBlank()) { - throw new IllegalArgumentException("prefix must not be blank"); - } + this.meterRegistry = Objects.requireNonNull(meterRegistry, "meterRegistry must not be null"); + this.prefix = Objects.requireNonNull(prefix, "prefix must not be null"); + if (prefix.isBlank()) { + throw new IllegalArgumentException("prefix must not be blank"); } + } - protected String validateName(String name){ - Objects.requireNonNull(name, "name must not be null"); - if(name.isBlank()){ - throw new IllegalArgumentException("name must not be blank"); - } - - return name.charAt(0) == '.' ? name : "." + name; + protected String validateName(String name) { + Objects.requireNonNull(name, "name must not be null"); + if (name.isBlank()) { + throw new IllegalArgumentException("name must not be blank"); } - protected String fullName(String name){ - return prefix + validateName(name); - } + return name.charAt(0) == '.' ? name : "." + name; + } - protected Counter newCounter(String name) { - return meterRegistry.counter(fullName(name)); - } + protected String fullName(String name) { + return prefix + validateName(name); + } - protected Gauge newGauge(String name, Supplier func) { - // no null checks in the builder below - Objects.requireNonNull(func, "func must not be null"); + protected Counter newCounter(String name) { + return meterRegistry.counter(fullName(name)); + } - return Gauge.builder(fullName(name), func) - .register(meterRegistry); - } + protected Gauge newGauge(String name, Supplier func) { + // no null checks in the builder below + Objects.requireNonNull(func, "func must not be null"); - protected Timer newTimer(String name){ - return newTimer(name, 0.5, 0.95, 0.99); - } + return Gauge.builder(fullName(name), func).register(meterRegistry); + } - protected Timer newTimer(String name, double... percentiles) { - return Timer.builder(fullName(name)) - .publishPercentiles(percentiles) - .register(meterRegistry); - } + protected Timer newTimer(String name) { + return newTimer(name, 0.5, 0.95, 0.99); + } + protected Timer newTimer(String name, double... percentiles) { + return Timer.builder(fullName(name)).publishPercentiles(percentiles).register(meterRegistry); + } - protected TimeGauge newTimeGauge(String name, Supplier func, TimeUnit unit) { + protected TimeGauge newTimeGauge(String name, Supplier func, TimeUnit unit) { - return TimeGauge.builder(fullName(name), func, unit) - .register(meterRegistry ); - } + return TimeGauge.builder(fullName(name), func, unit).register(meterRegistry); + } } diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/billing/AsyncBatchedLogUploader.java b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/AsyncBatchedLogUploader.java index dea042c9d0..b08631966a 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/service/billing/AsyncBatchedLogUploader.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/AsyncBatchedLogUploader.java @@ -2,32 +2,27 @@ import io.smallrye.mutiny.Uni; -/** - * A function that uploads a batch of log records, normally to S3. - */ +/** A function that uploads a batch of log records, normally to S3. */ @FunctionalInterface public interface AsyncBatchedLogUploader extends AutoCloseable { - /** - * Called to upload the batch of records. - * - * @param batch The batch of log records to upload - * @return A Uni of the result of the operation - */ - Uni upload(BatchedLogBuffer.Batch batch); + /** + * Called to upload the batch of records. + * + * @param batch The batch of log records to upload + * @return A Uni of the result of the operation + */ + Uni upload(BatchedLogBuffer.Batch batch); - @Override - default void close() {} + @Override + default void close() {} - /** - * Result of the upload call. - * @param success true if the operation succeeded, false otherwise. - * @param throwable The throwable associated with an error state. - * @param batch The batch that was uploaded, or attempted to be uploaded. - */ - record UploadResult( - boolean success, - Throwable throwable, - BatchedLogBuffer.Batch batch - ){} + /** + * Result of the upload call. + * + * @param success true if the operation succeeded, false otherwise. + * @param throwable The throwable associated with an error state. + * @param batch The batch that was uploaded, or attempted to be uploaded. + */ + record UploadResult(boolean success, Throwable throwable, BatchedLogBuffer.Batch batch) {} } diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BatchedLogBuffer.java b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BatchedLogBuffer.java index ee1efa3a58..eed43a3174 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BatchedLogBuffer.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BatchedLogBuffer.java @@ -1,12 +1,12 @@ package io.stargate.sgv2.jsonapi.service.billing; +import static io.stargate.sgv2.jsonapi.util.ClassUtils.classSimpleName; + import com.fasterxml.uuid.Generators; import com.fasterxml.uuid.NoArgGenerator; +import com.google.common.annotations.VisibleForTesting; import io.stargate.sgv2.jsonapi.metrics.BatchedLogBufferMetrics; -import io.stargate.sgv2.jsonapi.metrics.BatchedLogUploaderMetrics; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; - +import java.time.Clock; import java.time.Duration; import java.time.Instant; import java.util.*; @@ -14,87 +14,122 @@ import java.util.concurrent.BlockingQueue; import java.util.concurrent.atomic.AtomicLong; import java.util.logging.LogRecord; - -import static io.stargate.sgv2.jsonapi.util.ClassUtils.classSimpleName; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; /** - * Buffer for {@link LogRecord} that batches them according to the configuration. - *

- * See {@link #BatchedLogBuffer(int, long, Duration, int, BatchedLogUploaderMetrics)} for the config. - *

- *

- * There are two uses of this class, producers and consumers. + * Buffer for {@link LogRecord} that batches them according to the configuration when created. + * + *

There are two users of this class: + * *

    - *
  • Producers - call {@link #offer(LogRecord)} to add the log record to the buffer.
    • - *
  • Consumers - call {@link #nextBatch(boolean)} to get the next batch to send if there is a - * full batch.
    • + *
  • Producers - call {@link #offer(LogRecord)} to add the log record to the buffer. There can + * be many consumers form different threads. + *
  • Consumers - call {@link #nextBatch(boolean)} to get the next batch if available. There + * should be only 1 consumer calling at a time, caller is responsible for this. *
* * The buffer is designed to handle these as concurrent calls from different threads, and tracks * metrics for its use. - *

*/ public class BatchedLogBuffer { private static final Logger LOGGER = LoggerFactory.getLogger(BatchedLogBuffer.class); + @VisibleForTesting static final Clock DEFAULT_CLOCK = Clock.systemUTC(); + private final int maxBatchSize; - private final long maxBytes; - private final Duration maxAge; + private final long maxBatchBytes; + private final Duration maxBatchAge; + + // don't really need to keep this, just here for debugging + private final int capacity; + + // Clock we use to get the current time when checking age. + // so it can be controlled when testing + private final Clock clock; private final BlockingQueue queue; + private final AtomicLong queuedBytes = new AtomicLong(0); private final BatchedLogBufferMetrics metrics; - private final AtomicLong queuedBytes = new AtomicLong(0); + /** See {@link #BatchedLogBuffer(int, long, Duration, int, BatchedLogBufferMetrics, Clock)} */ + BatchedLogBuffer( + int maxBatchSize, + long maxBatchBytes, + Duration maxBatchAge, + int capacity, + BatchedLogBufferMetrics metrics) { + this(maxBatchSize, maxBatchBytes, maxBatchAge, capacity, metrics, DEFAULT_CLOCK); + } /** - * Creates a new instance of the buffer. + * Creates a new instance of the buffer with configuration. * - * @param maxBatchSize Maximum number of log records in a batch, when the buffer has more than this many - * entries a new batch is made available which will contain no more than this many lines. - * @param maxBatchBytes Maximum numbers of bytes in a batch, when the buffer has more than this many entries - * a new batch is made available which may contain more than this many bytes. The batch will - * have many maxBatchBytes if there is a single log record that is bigger. - * @param maxAge Maximum age any log record should have in the buffer before a new batch is available. - * @param queueCapacity Total number of log records to buffer. + *

NOTE: this ctor is for use in testing when the Clock is overridden, use the other + * ctor in regular code. + * + * @param maxBatchSize Maximum number of log records in a batch, when the buffer has more than + * this many entries a new batch is made available which will contain no more than this many + * lines. Batch maybe smaller if maxBatchBytes is hit. + * @param maxBatchBytes Maximum numbers of bytes in a batch, when the buffer has more than this + * many entries a new batch is made available which may contain more than this many bytes. The + * batch will have many maxBatchBytes if there is a single log record that is bigger. + * @param maxBatchAge Maximum age any log record should have in the buffer before a new batch is + * available. When a batch is triggered from max age the batch is filled, even if the other + * messages have not reached their max age. + * @param capacity Total number of log records to buffer. Beyond this called to {@link + * #offer(LogRecord)} will fail to add the message. * @param metrics Metrics recording object. + * @param clock The {@link Clock} implementation to use when checking the age of a message, this + * should only be overridden in testing. DO NOT USE IN CODE. If null uses {@link + * #DEFAULT_CLOCK} */ - public BatchedLogBuffer(int maxBatchSize, long maxBatchBytes, Duration maxAge, int queueCapacity, BatchedLogBufferMetrics metrics) { + @VisibleForTesting + BatchedLogBuffer( + int maxBatchSize, + long maxBatchBytes, + Duration maxBatchAge, + int capacity, + BatchedLogBufferMetrics metrics, + Clock clock) { if (maxBatchSize < 1) { - throw new IllegalArgumentException("maxBatchSize must be >= 1, got: " + maxBatchSize ); + throw new IllegalArgumentException("maxBatchSize must be >= 1, got: " + maxBatchSize); } if (maxBatchBytes < 1) { - throw new IllegalArgumentException("maxBatchBytes must be >= 1, got: " + maxBatchBytes ); - } - if (queueCapacity < 1) { - throw new IllegalArgumentException("queueCapacity must be >= 1, got: " + queueCapacity); + throw new IllegalArgumentException("maxBatchBytes must be >= 1, got: " + maxBatchBytes); } - if (maxAge == null || maxAge.isNegative() || maxAge.isZero()) { - throw new IllegalArgumentException("maxAge must be positive, got: " + maxAge); + if (maxBatchAge == null || maxBatchAge.isNegative() || maxBatchAge.isZero()) { + throw new IllegalArgumentException("maxAge must be positive, got: " + maxBatchAge); } + this.maxBatchSize = maxBatchSize; - this.maxBytes = maxBatchBytes; - this.maxAge = maxAge; + this.maxBatchBytes = maxBatchBytes; + this.maxBatchAge = maxBatchAge; + this.capacity = capacity; - this.metrics = Objects.requireNonNull(metrics, "billingMetrics must not be null"); + this.metrics = Objects.requireNonNull(metrics, "billingMetrics must not be null"); this.metrics.registerBuffer(this); + + this.clock = clock == null ? DEFAULT_CLOCK : clock; + if (this.clock != DEFAULT_CLOCK) { + LOGGER.warn( + "BatchedLogBuffer - WARNING - CONFIGURED TO USE A CUSTOM CLOCK, DO NOT USE IN PRODUCTION."); + } // must be concurrent to handle multiple threads - this.queue = new ArrayBlockingQueue<>(queueCapacity); + this.queue = new ArrayBlockingQueue<>(capacity); } /** * Appends the LogRecord to the buffer if the buffer has capacity. * - *

- * NOTE: because this is used for billing information if the record is - * null or has an empty message an exception is thrown rather than - * silently dropping it. We expect this situation to be an exception and it should fail. - *

- * @param record {@link LogRecord} to add to the buffer. - * @return true if the record was added to be buffer, false if the buffer did not have capacity. NOTE: - * this is different to the param check for record, the buffer filling is tracked as metric but no error. + *

NOTE: because this is used for billing information if the record is null or has an + * empty message an exception is thrown rather than silently dropping it. We expect this situation + * to be an exception and it should fail. * + * @param record {@link LogRecord} to add to the buffer. + * @return true if the record was added to be buffer, false if the buffer did not have capacity. */ public boolean offer(LogRecord record) { @@ -106,11 +141,11 @@ public boolean offer(LogRecord record) { } var newEntry = new Entry(record.getInstant(), logLine); - metrics.offered.increment(); + metrics.offered(); if (!queue.offer(newEntry)) { - // Bounded buffer full: drop and count + // Bounded buffer full, drop and count LOGGER.debug("offer() - buffer full, dropping new entry: {}", newEntry); - metrics.dropped.increment(); + metrics.dropped(); return false; } @@ -119,18 +154,17 @@ public boolean offer(LogRecord record) { } /** - * Returns the next batch of messages from the {@link LogRecord}'s added to the buffer, - * if one is available. - *

- * Designed to be called from different threads than those producing LogRecord's. - *

+ * Returns the next batch of messages from the {@link LogRecord}'s added to the buffer, if one is + * available. + * + *

Designed to be called from different threads than the producers called {@link + * #offer(LogRecord)} * - * @param drainFully when True a new batch is created without checking the configured - * rules, use this when draining the buffer and there may only - * be a partial batch. - * @return A new {@link Batch} of log messages all of which have been removed from the - * buffer, or null if there is no next batch. - * */ + * @param drainFully when True a new batch is created without checking the configured rules, use + * this when draining the buffer and there may only be a partial batch. + * @return A new {@link Batch} of log messages all of which have been removed from the buffer, or + * null if there is no next batch. + */ public Batch nextBatch(boolean drainFully) { var batchReason = decideNextBatch(drainFully); @@ -138,33 +172,63 @@ public Batch nextBatch(boolean drainFully) { return null; } - List lines = new ArrayList<>(maxBatchSize); + List batchLines = new ArrayList<>(maxBatchSize); Instant oldestEventAt = null; long batchBytes = 0; - Entry entry; + Entry peeked; // No matter why we started we create a full batch, e.g. we could start because the oldest // entry is past maxAge, but we still fill the batch. - while (lines.size() < maxBatchSize && batchBytes < maxBytes && (entry = queue.poll()) != null) { - - if (oldestEventAt == null || entry.eventAt().isBefore(oldestEventAt)) { - oldestEventAt = entry.eventAt(); + while (batchLines.size() < maxBatchSize && ((peeked = queue.peek()) != null)) { + + var lineBytes = peeked.lineBytes(); + if (batchBytes + lineBytes > maxBatchBytes && !batchLines.isEmpty()) { + // adding the next line will be too many bytes, we can only do this if the batch + // is empty, so a single big message can be put into a batch and not block everyone else + // break out of here. + break; } - lines.add(entry.line()); - var lineBytes = entry.lineBytes(); + // OK to remove entry from buffer and add to batch + // there is only this thread as a consumer, no race condition + var polled = queue.poll(); + // sanity check + if (polled != peeked) { + throw new IllegalStateException( + "nextBatch() - peeked entry is not same object as polled entry"); + } + if (oldestEventAt == null || polled.eventAt().isBefore(oldestEventAt)) { + oldestEventAt = polled.eventAt(); + } + batchLines.add(polled.line()); queuedBytes.addAndGet(-lineBytes); batchBytes += lineBytes; } - // sanity check, in case of concurrent calls - if (lines.isEmpty()){ + if (batchLines.isEmpty() && !queue.isEmpty()) { + // sanity check + // there is messages in the queue, but we did not put any in the batch, something wrong + // but it could be a race condition - things may be added after the loop finish + // so do not throw, just log + LOGGER.warn( + "nextBatch() - did not add any lines for next batch, but queue is not empty. May be logic bug or expected race condition. queue.size:{}", + queue.size()); return null; } - LOGGER.info("nextBatch() - next batch of log buffer, reason:{}, lines.size:{}, batchBytes:{}, oldestEventAt: {}", - batchReason, lines.size(), batchBytes, oldestEventAt); - return new Batch(batchReason, lines, batchBytes, oldestEventAt); + LOGGER.info( + "nextBatch() - next batch created, reason:{}, batchLines.size:{}, batchBytes:{}, oldestEventAt: {}", + batchReason, + batchLines.size(), + batchBytes, + oldestEventAt); + return new Batch(batchReason, batchLines, batchBytes, oldestEventAt, clock); + } + + /** Gets a copy of the contents of the buffer in a new array list, for testing. */ + @VisibleForTesting + List peekBuffer() { + return new ArrayList<>(queue); } public boolean isEmpty() { @@ -183,63 +247,112 @@ public int remainingCapacity() { return queue.remainingCapacity(); } + /** + * Gets the age of the item at the head of the buffer. + * + *

Age is determined by the clock used to create the buffer. + * + * @return age of the head item in the buffer, or null if no items in the buffer. + */ public Duration headEntryAge() { - var head = queue.peek(); - return head == null ? Duration.ZERO : Duration.between(head.eventAt(), Instant.now()); + return entryAge(queue.peek()); + } + + @VisibleForTesting + Duration entryAge(Entry entry) { + return entry == null ? Duration.ZERO : Duration.between(entry.eventAt(), clock.instant()); } private BillingBatchReason decideNextBatch(boolean drainFully) { - if (drainFully){ - return BillingBatchReason.DRAINING; - } - if (queue.size() >= maxBatchSize) { - return BillingBatchReason.MAX_BATCH_SIZE_EXCEEDED; - } - if (queuedBytes.get() > maxBytes) { - return BillingBatchReason.MAX_BYTES_EXCEEDED; + BillingBatchReason decision; + if (queue.isEmpty()) { + decision = null; + } else if (drainFully) { + decision = BillingBatchReason.DRAINING; + } else if (queue.size() >= maxBatchSize) { + decision = BillingBatchReason.MAX_SIZE_EXCEEDED; + } else if (queuedBytes.get() >= maxBatchBytes) { + decision = BillingBatchReason.MAX_BYTES_EXCEEDED; + } else if (headEntryAge().compareTo(maxBatchAge) >= 0) { + decision = BillingBatchReason.MAX_AGE_EXCEEDED; + } else { + decision = null; } - if (headEntryAge().compareTo(maxAge) >= 0){ - return BillingBatchReason.MAX_AGE_EXCEEDED; + if (LOGGER.isTraceEnabled()) { + LOGGER.trace("decideNextBatch() - drainFully:{} , decision:{}", drainFully, decision); } - return null; + return decision; } + @Override + public String toString() { + return new StringBuilder(classSimpleName(this) + "{") + .append("maxBatchSize=") + .append(maxBatchSize) + .append(", maxBatchBytes=") + .append(maxBatchBytes) + .append(", maxBatchAge=") + .append(maxBatchAge) + .append(", size=") + .append(size()) + .append("}") + .toString(); + } + + /** + * The reason a batch was created by the buffer. + * + *

... + */ public enum BillingBatchReason { DRAINING, - MAX_BATCH_SIZE_EXCEEDED, + MAX_SIZE_EXCEEDED, MAX_BYTES_EXCEEDED, MAX_AGE_EXCEEDED } /** + * A batch of log messages created by the buffer. * + *

See {@link BatchedLogBuffer#nextBatch(boolean)} */ public static final class Batch { private static final NoArgGenerator UUID_V7_GENERATOR = Generators.timeBasedEpochGenerator(); - private final UUID batchId = UUID_V7_GENERATOR.generate(); - private final BillingBatchReason batchReason; + private final UUID id = UUID_V7_GENERATOR.generate(); + private final BillingBatchReason reason; private final List lines; - private final long batchBytes; + private final long bytes; private final Instant oldestEventAt; - - public Batch(BillingBatchReason batchReason, List lines, long batchBytes, Instant oldestEventAt) { - this.batchReason = batchReason; - this.lines = Collections.unmodifiableList(lines); - this.batchBytes = batchBytes; - this.oldestEventAt = oldestEventAt; + private final Clock clock; + + Batch( + BillingBatchReason reason, + List lines, + long bytes, + Instant oldestEventAt, + Clock clock) { + this.reason = Objects.requireNonNull(reason, "reason must not be null"); + this.lines = + Collections.unmodifiableList(Objects.requireNonNull(lines, "lines must not be null")); + this.bytes = bytes; + this.oldestEventAt = Objects.requireNonNull(oldestEventAt, "oldestEventAt must not be null"); + this.clock = Objects.requireNonNull(clock, "clock must not be null"); } public UUID id() { - return batchId; + return id; } public BillingBatchReason reason() { - return batchReason; + return reason; } + /** + * @return Unmodifiable list of the log lines in this buffer + */ public List lines() { return lines; } @@ -249,39 +362,64 @@ public Instant oldestEventAt() { } public Duration oldestEventAtDuration() { - return Duration.between(oldestEventAt(), Instant.now()); + // using the outer buffers clock, so any tests that change the clock + // get a consistent results + return Duration.between(oldestEventAt(), clock.instant()); } public int size() { return lines.size(); } - public long batchBytes() { - return batchBytes; + public long bytes() { + return bytes; } - public String description() { - return "id:%s, reason:%s, oldestEventAt:%s, size:%s, batchBytes:%s".formatted( - batchId, batchReason, oldestEventAt, size(), batchBytes - ); + @Override + public String toString() { + return new StringBuilder(classSimpleName(this) + "{") + .append("id=") + .append(id) + .append(", reason=") + .append(reason) + .append(", oldestEventAt=") + .append(oldestEventAt) + .append(", size=") + .append(size()) + .append(", bytes=") + .append(bytes()) + .append("}") + .toString(); } } /** - * Holder for the billing event lines we get called with. - * @param eventAt When the event happened - * @param line The billing event line to record + * An entry in a batch, this is one log message passed to the buffer. + * + * @param eventAt When the log event happened + * @param line The log message */ - private record Entry(Instant eventAt, String line) { + public record Entry(Instant eventAt, String line) { /** * Gets the length of the line in bytes, - *

- * Kind of a hack, we are counting unicode code points and calling that 1 byte. Should work - * for ascii text, will undercount if there is non ascii chars but everthing in billing should be ascii + * + *

Kind of a hack, we are counting Unicode code points and calling that 1 byte. Should work + * for ascii text, will undercount if there is non ASCII chars but everything in billing should + * be ascii + * + * @return length of the line in bytes, included a carriage return for `\n` */ public int lineBytes() { - return line.length() + 1; // +1 is for a newline + return lineBytes(line); + } + + /** + * @return length of the line in bytes, included a carriage return for `\n` + */ + @VisibleForTesting + static int lineBytes(String line) { + return line.length() + 1; } } } diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3HandlerInstaller.java b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3HandlerInstaller.java index 384fc8796f..b975262e04 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3HandlerInstaller.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3HandlerInstaller.java @@ -13,18 +13,17 @@ import org.slf4j.LoggerFactory; /** - * - *

Attaches a {@link BillingS3LogHandler} to the {@code billing.events} logger at startup - * (when {@link BillingS3ExportConfig#enabled()} is {@code true}) and removes + closes it on - * shutdown for a graceful drain. - * + * Attaches a {@link BillingS3LogHandler} to the {@code billing.events} logger at startup (when + * {@link BillingS3ExportConfig#enabled()} is {@code true}) and removes + closes it on shutdown for + * a graceful drain. */ @ApplicationScoped public class BillingS3HandlerInstaller { - private static final org.slf4j.Logger LOGGER = LoggerFactory.getLogger(BillingS3HandlerInstaller.class); + private static final org.slf4j.Logger LOGGER = + LoggerFactory.getLogger(BillingS3HandlerInstaller.class); - private static final String METRICS_PREFIX ="billing"; + private static final String METRICS_PREFIX = "billing"; private static final String BILLING_LOGGER_NAME = "billing.events"; private final BillingS3ExportConfig config; @@ -44,32 +43,33 @@ void onStart(@Observes StartupEvent event) { LOGGER.info("Billing S3 export disabled"); return; } + LOGGER.info("Billing S3 export enabled"); // Fail-loud: invalid billing S3 config throws here, aborting application startup. - var uploader = S3BatchedLogUploader.create( + var uploader = + S3BatchedLogUploader.create( config.region(), config.bucket(), config.endpointOverride().orElse(null), new BatchedLogUploaderMetrics(meterRegistry, METRICS_PREFIX)); + LOGGER.info("Billing is using uploader: {}", uploader); - var buffer= new BatchedLogBuffer( + var buffer = + new BatchedLogBuffer( config.maxEventsPerBatch(), config.maxBytesPerBatch(), config.maxAge(), config.queueCapacity(), new BatchedLogBufferMetrics(meterRegistry, METRICS_PREFIX)); - - this.handler = new BillingS3LogHandler(config, uploader, meterRegistry); + LOGGER.info("Billing is using log buffer: {}", buffer); + this.handler = new BillingS3LogHandler(buffer, uploader); // TODO: LOGGER NAME SHOULD BE IN CONFIG Logger.getLogger(BILLING_LOGGER_NAME).addHandler(this.handler); - LOGGER.info( - "Attached billing S3 export handler to logger named {}, bucket={}, region={}, endpointOverride={}", - BILLING_LOGGER_NAME, - bucket, - region, - config.endpointOverride().orElse(null)); + "Billing has attached BillingS3LogHandler to the logger named: {}", BILLING_LOGGER_NAME); + + // TODO: XXXX call start on the thread. } void onStop(@Observes ShutdownEvent event) { diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3LogHandler.java b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3LogHandler.java index 58c53225eb..e1058524bc 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3LogHandler.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3LogHandler.java @@ -2,7 +2,6 @@ import com.google.common.annotations.VisibleForTesting; import io.smallrye.mutiny.Uni; - import java.time.Duration; import java.util.concurrent.CountDownLatch; import java.util.concurrent.TimeUnit; @@ -19,9 +18,10 @@ *

See {@link BillingS3HandlerInstaller} for setup. // AI SLOP BELOW JUL handler that turns * {@code billing.events} log lines into batched S3 objects. * - *

Division of labor: {@link BatchedLogBuffer} decides when a batch seals, {@link AsyncBatchedLogUploader} - * decides what an S3 object looks like, and this class decides when uploads run — the flush - * triggers (seal on publish, age tick, drain on close), the upload-concurrency gate, and metrics. + *

Division of labor: {@link BatchedLogBuffer} decides when a batch seals, {@link + * AsyncBatchedLogUploader} decides what an S3 object looks like, and this class decides when + * uploads run — the flush triggers (seal on publish, age tick, drain on close), the + * upload-concurrency gate, and metrics. * *

Delivery is at-most-once by design: publish never waits for queue capacity, full buffers drop * new lines, and close drains best-effort within {@code shutdownTimeout}. @@ -60,8 +60,7 @@ public final class BillingS3LogHandler extends Handler { private final BatchedLogBuffer batchedLogBuffer; @VisibleForTesting - BillingS3LogHandler( - AsyncBatchedLogUploader uploader, BatchedLogBuffer batchedLogBuffer) { + BillingS3LogHandler(BatchedLogBuffer batchedLogBuffer, AsyncBatchedLogUploader uploader) { this.batchedLogBuffer = batchedLogBuffer; this.uploader = uploader; diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/billing/S3BatchedLogUploader.java b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/S3BatchedLogUploader.java index 5eec6d1fd7..bcc93fe2ad 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/service/billing/S3BatchedLogUploader.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/S3BatchedLogUploader.java @@ -1,7 +1,9 @@ package io.stargate.sgv2.jsonapi.service.billing; -import com.google.common.annotations.VisibleForTesting; +import static io.stargate.sgv2.jsonapi.util.ClassUtils.classSimpleName; + import io.smallrye.mutiny.Uni; +import io.stargate.sgv2.jsonapi.metrics.BatchedLogUploaderMetrics; import java.net.URI; import java.nio.charset.StandardCharsets; import java.time.Duration; @@ -9,8 +11,6 @@ import java.time.format.DateTimeFormatter; import java.util.Objects; import java.util.concurrent.CompletionException; - -import io.stargate.sgv2.jsonapi.metrics.BatchedLogUploaderMetrics; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import software.amazon.awssdk.awscore.exception.AwsServiceException; @@ -19,11 +19,11 @@ import software.amazon.awssdk.services.s3.S3AsyncClient; import software.amazon.awssdk.services.s3.model.PutObjectRequest; -/** Uploads sealed billing batches to S3 as NDJSON objects under time-partitioned keys. - * TODO: - * .requestChecksumCalculation(RequestChecksumCalculation.WHEN_SUPPORTED) - * .responseChecksumValidation(ResponseChecksumValidation.WHEN_SUPPORTED) - * */ +/** + * Uploads sealed billing batches to S3 as NDJSON objects under time-partitioned keys. TODO: + * .requestChecksumCalculation(RequestChecksumCalculation.WHEN_SUPPORTED) + * .responseChecksumValidation(ResponseChecksumValidation.WHEN_SUPPORTED) + */ public class S3BatchedLogUploader implements AsyncBatchedLogUploader { private static final Logger LOGGER = LoggerFactory.getLogger(S3BatchedLogUploader.class); @@ -34,7 +34,6 @@ public class S3BatchedLogUploader implements AsyncBatchedLogUploader { private static final DateTimeFormatter OBJECT_KEY_FORMATTER = DateTimeFormatter.ofPattern("yyyy/MM/dd/HH/mm").withZone(ZoneOffset.UTC); - private static final Duration API_CALL_ATTEMPT_TIMEOUT = Duration.ofSeconds(10); private static final Duration API_CALL_TIMEOUT = Duration.ofSeconds(30); @@ -43,11 +42,8 @@ public class S3BatchedLogUploader implements AsyncBatchedLogUploader { private final String bucket; private final BatchedLogUploaderMetrics metrics; - - private S3BatchedLogUploader(S3AsyncClient client, - String region, - String bucket, - BatchedLogUploaderMetrics metrics) { + private S3BatchedLogUploader( + S3AsyncClient client, String region, String bucket, BatchedLogUploaderMetrics metrics) { this.client = client; this.region = region; this.bucket = bucket; @@ -62,10 +58,8 @@ private S3BatchedLogUploader(S3AsyncClient client, * @param endpointOverride * @return */ - public static S3BatchedLogUploader create(String region, - String bucket, - String endpointOverride, - BatchedLogUploaderMetrics metrics) { + public static S3BatchedLogUploader create( + String region, String bucket, String endpointOverride, BatchedLogUploaderMetrics metrics) { if (region == null || region.isBlank()) { throw new IllegalArgumentException("region must not be null or blank"); @@ -73,14 +67,15 @@ public static S3BatchedLogUploader create(String region, if (bucket == null || bucket.isBlank()) { throw new IllegalArgumentException("bucket must not be null or blank"); } - Objects.requireNonNull(metrics, "metrics must not be null"); + Objects.requireNonNull(metrics, "metrics must not be null"); // Credentials resolve from the SDK's default provider chain (env vars, web-identity/OIDC // token, instance/container roles), left implicit so the client owns — and closes — the // provider. This transparently supports federated (AssumeRoleWithWebIdentity) and // cross-account access: the bucket may live in a different account (per IAM + bucket // policy); its region is set via .region(). - var builder = S3AsyncClient.builder() + var builder = + S3AsyncClient.builder() .region(Region.of(region)) .overrideConfiguration( config -> @@ -93,10 +88,10 @@ public static S3BatchedLogUploader create(String region, // An override is only for a non-AWS S3 (S3Mock in tests): it bypasses those rules and forces // path-style, since a localhost host can't virtual-host the bucket as a subdomain. if (endpointOverride != null) { - builder.endpointOverride(URI.create(endpointOverride)).forcePathStyle(true); + builder.endpointOverride(URI.create(endpointOverride)).forcePathStyle(Boolean.TRUE); } - return new S3BatchedLogUploader(builder.build(), region, bucket, metrics); + return new S3BatchedLogUploader(builder.build(), region, bucket, metrics); } @Override @@ -107,19 +102,19 @@ public Uni upload(BatchedLogBuffer.Batch batch) { var location = objectLocation(batch); var body = objectContent(batch); - LOGGER.info( "upload() - starting to upload batch, batch:({}), location:{}, body.size:{}", - batch.description(), + batch, location, - body.length); + String.valueOf(body.length)); // No .retry() here: unconfigured, S3AsyncClient already retries (default LegacyRetryStrategy — // 3 retries / 4 attempts). See // https://docs.aws.amazon.com/sdk-for-java/latest/developer-guide/retry-strategy.html // and see https://github.com/aws/aws-sdk-java-v2/issues/6987 for future change. - var putRequest = PutObjectRequest.builder() + var putRequest = + PutObjectRequest.builder() .bucket(location.bucket()) .key(location.key()) .contentType(CONTENT_TYPE_NDJSON) @@ -138,7 +133,7 @@ public Uni upload(BatchedLogBuffer.Batch batch) { metrics.recordBatchFailed(batch); LOGGER.error( "upload() - error uploading billing to S3, batch:({}), location:{}, requestId:{}", - batch.description(), + batch, location, requestId, cause); @@ -146,13 +141,13 @@ public Uni upload(BatchedLogBuffer.Batch batch) { metrics.recordBatchDelivered(batch); LOGGER.info( "upload() - success uploading billing to S3, batch:({}), location:{}, requestId:{}, eTag:{}, status:{}", - batch.description(), + batch, location, requestId, resp.eTag(), - resp.sdkHttpResponse().statusCode()); + String.valueOf(resp.sdkHttpResponse().statusCode())); } - return new UploadResult(success, failure, batch); + return new UploadResult(success, failure, batch); }); } @@ -161,7 +156,19 @@ public void close() { client.close(); } - private S3Location objectLocation(BatchedLogBuffer.Batch batch) { + @Override + public String toString() { + return new StringBuilder(classSimpleName(this) + "{") + .append("region=") + .append(region) + .append(", bucket=") + .append(bucket) + .append(", pathPrefix=") + .append(PATH_PREFIX) + .toString(); + } + + private S3Location objectLocation(BatchedLogBuffer.Batch batch) { var objectKey = PATH_PREFIX @@ -183,8 +190,5 @@ private static byte[] objectContent(BatchedLogBuffer.Batch batch) { return sb.toString().getBytes(StandardCharsets.UTF_8); } - private record S3Location( - String region, - String bucket, - String key) {} + private record S3Location(String region, String bucket, String key) {} } diff --git a/src/test/java/io/stargate/sgv2/jsonapi/TestConstants.java b/src/test/java/io/stargate/sgv2/jsonapi/TestConstants.java index 35be0bddc5..8dd3949a92 100644 --- a/src/test/java/io/stargate/sgv2/jsonapi/TestConstants.java +++ b/src/test/java/io/stargate/sgv2/jsonapi/TestConstants.java @@ -18,11 +18,11 @@ import io.stargate.sgv2.jsonapi.config.constants.DocumentConstants; import io.stargate.sgv2.jsonapi.config.feature.ApiFeatures; import io.stargate.sgv2.jsonapi.metrics.JsonProcessingMetricsReporter; +import io.stargate.sgv2.jsonapi.service.billing.Billing; import io.stargate.sgv2.jsonapi.service.cqldriver.CQLSessionCache; import io.stargate.sgv2.jsonapi.service.cqldriver.executor.*; import io.stargate.sgv2.jsonapi.service.embedding.operation.EmbeddingProvider; import io.stargate.sgv2.jsonapi.service.embedding.operation.EmbeddingProviderFactory; -import io.stargate.sgv2.jsonapi.service.billing.Billing; import io.stargate.sgv2.jsonapi.service.reranking.operation.RerankingProviderFactory; import io.stargate.sgv2.jsonapi.service.schema.*; import io.stargate.sgv2.jsonapi.service.schema.collections.CollectionLexicalDefSchemaFactory; diff --git a/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BatchedLogBufferTest.java b/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BatchedLogBufferTest.java new file mode 100644 index 0000000000..7baf16ffa2 --- /dev/null +++ b/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BatchedLogBufferTest.java @@ -0,0 +1,913 @@ +package io.stargate.sgv2.jsonapi.service.billing; + +import static io.stargate.sgv2.jsonapi.util.ClassUtils.classSimpleName; +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; +import static org.awaitility.Awaitility.await; +import static org.mockito.Mockito.*; + +import io.stargate.sgv2.jsonapi.metrics.BatchedLogBufferMetrics; +import io.stargate.sgv2.jsonapi.util.MockClock; +import java.lang.ref.WeakReference; +import java.time.Duration; +import java.util.*; +import java.util.concurrent.*; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicLong; +import java.util.concurrent.locks.LockSupport; +import java.util.logging.Level; +import java.util.logging.LogRecord; +import java.util.stream.IntStream; +import java.util.stream.Stream; +import org.junit.Test; + +/** + * Unit tests for {@link BatchedLogBuffer}\ + * + *

TODO: out of order log records gets correct oldest metric TODO: TEST a big line bigger than + * the max bytes gets through TODO: test metrics using SimpleMeterRegistry + */ +public class BatchedLogBufferTest { + + // want the line bytes when lines go into the buffer to be 25 + // template below is 21 bytes + // 3 chars for the index get added in createFixture() + // 1 char added in the buffer calc's for the `\n` to write out + private static final int MESSAGE_LENGTH_IN_BUFFER = 25; + private static final String TEMPLATE_25_CHARS = "Total of 25 chars "; + + private static final int MAX_BATCH_SIZE = 100; + // The number of messages we can fit inside the max bytes setting + private static final int MAX_BATCH_BYTES_NUM_MESSAGES = 20; + private static final int MAX_BATCH_BYTES = + MESSAGE_LENGTH_IN_BUFFER * MAX_BATCH_BYTES_NUM_MESSAGES; + + // How many full batches, tracked by max size, we want to fit in the buffer + private static final int BATCHES_BY_SIZE_PER_CAPACITY = 3; + private static final int BUFFER_CAPACITY = MAX_BATCH_SIZE * BATCHES_BY_SIZE_PER_CAPACITY; + + // number of log records we create for each feature / test + private static final int NUM_RECORDS = BUFFER_CAPACITY * 3; + // when using mock clock, we set the instant for each log record to be 1 "second" + // after the last, so we will create log records with up to + // NUM_RECORDS of seconds past when the clock was started + // used when testing the max age features + private static final Duration MAX_AGE = Duration.ofSeconds(NUM_RECORDS); + private static final Level LOG_LEVEL = Level.INFO; + + // ********************************************************* + // Offer - Producer side of the buffer + // ********************************************************* + + /** When the buffer reaches capacity calling offer() fails. Single producer thread. */ + @Test + public void offerFailsAtCapacitySingleThread() { + + var fixture = defaultFixture(false); + var snapshot = BufferSnapshot.create(fixture); + var slice = Slice.to(BUFFER_CAPACITY); + + // send full capacity to the buffer, should all work + fixture.assertOffer("offerFailsAtCapacitySingleThread() - prefill to capacity", slice); + + // check the change in the buffer is expected given the slice of source data + snapshot.assertAll("offerFailsAtCapacitySingleThread()", slice, true); + // Buffer should now be full, try to add one more + fixture.assertBufferFull("offerFailsAtCapacitySingleThread()", BUFFER_CAPACITY + 1); + } + + /** When the buffer reaches capacity calling offer() fails. Multiple producer threads. */ + @Test + public void offerFailsAtCapacityMultiThread() { + + var fixture = defaultFixture(false); + var snapshot = BufferSnapshot.create(fixture); + var slice = Slice.to(BUFFER_CAPACITY); + + // fill the buffer to capacity from 6 threads calling offer() + // auto close will wait for tasks to finish in executor + try (var pool = Executors.newFixedThreadPool(6)) { + for (var record : slice.stream(fixture.logRecords).toList()) { + pool.submit(() -> fixture.buffer.offer(record)); + } + } + + // check the change in the buffer is expected given the slice of source data + snapshot.assertAll("offerFailsAtCapacityMultiThread()", slice, false); + // Buffer should now be full, try to add one more + fixture.assertBufferFull("offerFailsAtCapacityMultiThread()", BUFFER_CAPACITY + 1); + } + + /** + * Verify that when offered a LogRecord the buffer does not hold reference to the LogRecord and it + * can be GC'd + */ + @Test + public void offerDoesNotHoldReferences() { + + var fixture = defaultFixture(false); + + // do not use the records in the fixture, they are held in a list + var record = new LogRecord(Level.INFO, "offerDoesNotHoldReferences()"); + var ref = new WeakReference<>(record); + + fixture.buffer.offer(record); + record = null; + + // reference count for the object created for "record" above should now be zero + // will timeout if the object is not GC'd and error + await("offerDoesNotHoldReferences() - waiting for record to be GC'd") + .atMost(Duration.ofSeconds(5)) + .until( + () -> { + System.gc(); + return ref.get() == null; + }); + } + + @Test + public void offerNullRecord() { + var fixture = defaultFixture(false); + + assertThatThrownBy(() -> fixture.buffer.offer(null)) + .as("offerNullRecord() null log record is an exception") + .isInstanceOf(NullPointerException.class); + } + + @Test + public void offerNullOrBlankMessage() { + var fixture = defaultFixture(false); + + var nullRecord = new LogRecord(Level.INFO, null); + var blankRecord = new LogRecord(Level.INFO, " "); + + assertThatThrownBy(() -> fixture.buffer.offer(nullRecord)) + .as("offerNullOrBlankMessage() - null message is an error") + .isInstanceOf(IllegalArgumentException.class); + assertThatThrownBy(() -> fixture.buffer.offer(blankRecord)) + .as("offerNullOrBlankMessage() - blank message is an error") + .isInstanceOf(IllegalArgumentException.class); + } + + // ********************************************************* + // nextBatch - Consumer side of the buffer + // ********************************************************* + + /** When buffer is empty, there is no batch available. */ + @Test + public void nextBatchEmptyBufferNoBatch() { + + // lock the clock, do not want it to auto advance for batch testing + var fixture = defaultFixture(true); + + assertThat(fixture.buffer.nextBatch(false)) + .as("nextBatchEmptyBufferNoBatch() - drainFully=false, no batch") + .isNull(); + + assertThat(fixture.buffer.nextBatch(true)) + .as("nextBatchEmptyBufferNoBatch() - drainFully=true, no batch") + .isNull(); + } + + /** Properties of the returned batch object are as expected. */ + @Test + public void nextBatchBatchProperties() { + + // lock the clock, do not want it to auto advance for batch testing + var fixture = defaultFixture(true); + var slice = Slice.to(BUFFER_CAPACITY); + + // fill the buffer with all the records it will fit + fixture.assertOffer("nextBatchBatchProperties()", slice); + + // keep taking batches and check their properties + BatchedLogBuffer.Batch batch; + Set batchIds = new HashSet<>(); + while ((batch = fixture.buffer.nextBatch(true)) != null) { + + assertThat(batch.id()) + .as("nextBatchBatchProperties() - batch ID has not been seen") + .satisfies(batchIds::add); + + assertThat(batch.toString()) + .as("nextBatchBatchProperties() - batch toString has values") + .contains("id=" + batch.id()) + .contains("reason=" + batch.reason()) + .contains("size=" + batch.size()) + .contains("bytes=" + batch.bytes()); + } + + assertThat(fixture.buffer.isEmpty()) + .as("nextBatchBatchProperties() - drained buffer is empty") + .isTrue(); + } + + /** Metadata (size etc) for the buffer is updated after a batch is returned. */ + @Test + public void nextBatchMetaUpdatedAfterBatch() { + + // lock the clock, do not want it to auto advance for batch testing + var fixture = defaultFixture(true); + var slice = Slice.to(MAX_BATCH_SIZE); + + // fill the buffer with 1 batch size and assert metadata + fixture.assertOffer("nextBatchMetaUpdatedAfterBatch()", slice); + + // take 1 batch + // we are only checking that the bookkeeping on the buffer changes, not checking + // rules for batch selections, this is done by assertNextBatch() + var batch1 = fixture.assertNextBatch("nextBatchMetaUpdatedAfterBatch() - 1st", false); + + // take a second batch and check again bookkeeping updated + var batch2 = fixture.assertNextBatch("nextBatchMetaUpdatedAfterBatch() - 2nd", false); + } + + /** Trigger a batch from the number of records added to buffer */ + @Test + public void nextBatchTriggerMaxSize() { + + // change so the template is small so does not trigger max bytes + // lock the clock, do not want it to auto advance for batch testing + var fixture = + createFixture( + MAX_BATCH_SIZE, + MAX_BATCH_BYTES * 100, // big number so never batch because of bytes + MAX_AGE, + BUFFER_CAPACITY, + NUM_RECORDS, + LOG_LEVEL, + "test-", + true); + + // Fill to 1 less than max batch size, should be no batch + var slice1 = Slice.to(MAX_BATCH_SIZE - 1); + fixture.assertOffer("nextBatchMetaUpdatedAfterBatch()", slice1); + var batch1 = fixture.buffer.nextBatch(false); + assertThat(batch1).as("nextBatchTriggerMaxSize() - < MAX_BATCH_SIZE, no batch").isNull(); + + // add one more record, should be a batch of MAX_BATCH_SIZE + var slice2 = Slice.slice(MAX_BATCH_SIZE - 1, MAX_BATCH_SIZE); + fixture.assertOffer("nextBatchMetaUpdatedAfterBatch()", slice2); + var batch2 = fixture.assertNextBatch("nextBatchMetaUpdatedAfterBatch() - 2nd", false); + + assertThat(batch2.size()) + .as("nextBatchTriggerMaxSize() - 2nd batch is full batch size") + .isEqualTo(MAX_BATCH_SIZE); + + assertThat(batch2.reason()) + .as( + "nextBatchTriggerMaxSize() - 2nd batch because " + + BatchedLogBuffer.BillingBatchReason.MAX_SIZE_EXCEEDED) + .isEqualTo(BatchedLogBuffer.BillingBatchReason.MAX_SIZE_EXCEEDED); + + // add one more record, should be no more batches + var slice3 = Slice.slice(MAX_BATCH_SIZE, MAX_BATCH_SIZE + 1); + fixture.assertOffer("nextBatchMetaUpdatedAfterBatch() - 3rd", slice3); + var batch3 = fixture.buffer.nextBatch(false); + assertThat(batch3).as("nextBatchTriggerMaxSize() - 3rd - no batch").isNull(); + } + + /** Trigger a batch from the byte size in the buffer */ + @Test + public void nextBatchTriggerMaxBytes() { + + // default fixture will only fit + // the MAX_BATCH_BYTES_NUM_MESSAGES which is less than MAX_SIZE + // lock the clock, do not want it to auto advance for batch testing + var fixture = defaultFixture(true); + + // Fill to 1 message less than max bytes size, should be no batch + var slice1 = Slice.to(MAX_BATCH_BYTES_NUM_MESSAGES - 1); + fixture.assertOffer("nextBatchTriggerMaxBytes()", slice1); + var batch1 = fixture.buffer.nextBatch(false); + assertThat(batch1).as("nextBatchTriggerMaxSize() - < MAX_BATCH_BYTES, no batch").isNull(); + + // add one more , should be a batch of full batch bytes + var slice2 = Slice.slice(MAX_BATCH_BYTES_NUM_MESSAGES - 1, MAX_BATCH_BYTES_NUM_MESSAGES); + fixture.assertOffer("nextBatchTriggerMaxBytes()", slice2); + var batch2 = fixture.assertNextBatch("nextBatchTriggerMaxBytes() - 2nd", false); + + // we know how many we put in there + assertThat(batch2.bytes()) + .as("nextBatchTriggerMaxBytes() - 2nd batch byte size match") + .isEqualTo(MAX_BATCH_BYTES_NUM_MESSAGES * MESSAGE_LENGTH_IN_BUFFER); + + assertThat(batch2.reason()) + .as( + "nextBatchTriggerMaxBytes() - 2nd batch because " + + BatchedLogBuffer.BillingBatchReason.MAX_BYTES_EXCEEDED) + .isEqualTo(BatchedLogBuffer.BillingBatchReason.MAX_BYTES_EXCEEDED); + + // add one more, should be no more batches + var slice3 = Slice.slice(MAX_BATCH_BYTES_NUM_MESSAGES, MAX_BATCH_BYTES_NUM_MESSAGES + 1); + fixture.assertOffer("nextBatchTriggerMaxBytes() - 3rd", slice3); + var batch3 = fixture.buffer.nextBatch(false); + assertThat(batch3).as("nextBatchTriggerMaxBytes() - 3rd - no batch").isNull(); + } + + /** Trigger a batch from the maximum age of the first element in the buffer */ + @Test + public void nextBatchTriggerMaxAge() { + + // lock the clock, do not want it to auto advance for batch testing + // NOTE: WE ARE USING THE MOCK CLOCK IN THIS TEST, WE CONTROL TIME + var fixture = defaultFixture(true); + + // Add only 3 messages, we will not trip size or bytes tigger + final int ADDED_RECORDS = 3; + var slice1 = Slice.to(ADDED_RECORDS); + fixture.assertOffer("nextBatchTriggerMaxAge()", slice1); + + // the clock has not moved, there should be no batch + var batch1 = fixture.buffer.nextBatch(false); + assertThat(batch1).as("nextBatchTriggerMaxAge() - clock as not moved, no batch").isNull(); + + // Every LogRecord created in fixture has an instanceAt of 1 second after the previous + // the first LogRecord has the same instanceAt as when the clock started. + // so if we advance the clock to be MAX_AGE after when it started the only LogRecord that will + // be too old is the first, the others are all 1+ seconds younger + var newNow = fixture.clock().startedAt().plus(MAX_AGE); + fixture.clock().setInstant(newNow); + + // The buffer should now think the time is "newNow" + // Sanity check, before getting the batch check that only the first log record is MAX_AGE + // checking all this junk did what I think + int i = 0; + var peekedBuffer = fixture.buffer.peekBuffer(); + for (var peekEntry : peekedBuffer) { + var entryAge = fixture.buffer.entryAge(peekEntry); + if (i == 0) { + assertThat(entryAge) + .as("nextBatchTriggerMaxAge() - clock moved, first entry should be MAX_AGE old") + .isEqualTo(MAX_AGE); + } else { + assertThat(entryAge) + .as( + "nextBatchTriggerMaxAge() - clock moved, non first entry should be < MAX_AGE old. i: " + + i) + .isLessThan(MAX_AGE); + } + i++; + } + + // with the clock advanced the buffer should now trigger a batch because + // MAX_AGE_EXCEEDED + var batch2 = fixture.assertNextBatch("nextBatchTriggerMaxAge() - 2nd", false); + + assertThat(batch2.reason()) + .as( + "nextBatchTriggerMaxAge() - 2nd batch because " + + BatchedLogBuffer.BillingBatchReason.MAX_AGE_EXCEEDED) + .isEqualTo(BatchedLogBuffer.BillingBatchReason.MAX_AGE_EXCEEDED); + + // should have drained all the messages, even if they were not too old + assertThat(fixture.buffer.size()) + .as("nextBatchTriggerMaxAge() - 2nd batch buffer, size") + .isEqualTo(0); + assertThat(fixture.buffer.isEmpty()) + .as("nextBatchTriggerMaxAge() - 2nd batch buffer, isEmpty") + .isTrue(); + assertThat(fixture.buffer.queuedBytes()) + .as("nextBatchTriggerMaxAge() - 2nd batch buffer, bytes") + .isEqualTo(0); + + // sanity check, we should have ADDED_RECORDS entries in the batch + // and the oldest should be the first one we created + assertThat(batch2.size()) + .as("nextBatchTriggerMaxAge() - 2nd batch buffer, size expected") + .isEqualTo(ADDED_RECORDS); + assertThat(batch2.oldestEventAt()) + .as("nextBatchTriggerMaxAge() - 2nd batch buffer, oldest event expected") + .isEqualTo(fixture.logRecords.getFirst().getInstant()); + + // add one more record, should be no more batches + var slice3 = Slice.slice(ADDED_RECORDS, ADDED_RECORDS + 1); + fixture.assertOffer("nextBatchTriggerMaxAge() - 3rd", slice3); + var batch3 = fixture.buffer.nextBatch(false); + assertThat(batch3).as("nextBatchTriggerMaxAge() - 3rd - no batch").isNull(); + } + + /** + * Trigger a batch because drainFully=true so we want everything from it regardless of size, + * bytes, age + */ + @Test + public void nextBatchTriggerDrain() { + + // lock the clock, do not want it to auto advance for batch testing + var fixture = defaultFixture(true); + + // Fill so we have 1 full batch and 1 partial batch + var PARTIAL_BATCH_SIZE = 10; + var slice1 = Slice.to(MAX_BATCH_BYTES_NUM_MESSAGES + PARTIAL_BATCH_SIZE); + fixture.assertOffer("nextBatchTriggerDrain()", slice1); + + // 1st - drainFully - should get a full batch + var batch1 = fixture.assertNextBatch("nextBatchTriggerDrain() - 1st - full batch", true); + assertThat(batch1.reason()) + .as( + "nextBatchTriggerDrain() - 1st - reason is " + + BatchedLogBuffer.BillingBatchReason.DRAINING) + .isEqualTo(BatchedLogBuffer.BillingBatchReason.DRAINING); + assertThat(batch1.size()) + .as("nextBatchTriggerDrain() - 1st - full batch, size") + .isEqualTo(MAX_BATCH_BYTES_NUM_MESSAGES); + + // 2nd - drainFully - should get a partial batch + var batch2 = fixture.assertNextBatch("nextBatchTriggerDrain() - 2nd - partial batch", true); + assertThat(batch1.reason()) + .as( + "nextBatchTriggerDrain() - 2nd - reason is " + + BatchedLogBuffer.BillingBatchReason.DRAINING) + .isEqualTo(BatchedLogBuffer.BillingBatchReason.DRAINING); + assertThat(batch2.size()) + .as("nextBatchTriggerDrain() - 2nd - partial batch, size") + .isEqualTo(PARTIAL_BATCH_SIZE); + + // 3rs - drainFully - no more batch + var batch3 = fixture.buffer.nextBatch(true); + assertThat(batch3).as("nextBatchTriggerMaxBytes() - 3rd - no batch").isNull(); + } + + /** + * Multiple producers sending to the buffer, and one consumer reading from it concurrently. + * + *

NOTE: ttest takes 7 or 8 seconds, if you change the sleep time it may mean there are + * no batches collected after shutdown because producers go fast + */ + @Test + public void multiThreadedProducerConsumer() { + + var fixture = defaultFixture(false); + + // Setup a Consumer thread, it will keep running until we set consumerShutdown + var normalBatches = new ArrayList(); + var shutdownBatches = new ArrayList(); + var consumerShutdown = new AtomicBoolean(false); + var consumerExecutor = + Executors.newSingleThreadExecutor(Thread.ofPlatform().name("consumer-", 0).factory()); + + var consumerFuture = + consumerExecutor.submit( + () -> { + // this is the consumer in normal operations, read batches, if none sleep, read again + while (!consumerShutdown.get()) { + BatchedLogBuffer.Batch consumerNormalBatch; + // drainFully=false - because not trying to shutdown + while ((consumerNormalBatch = fixture.buffer.nextBatch(false)) != null) { + normalBatches.add(consumerNormalBatch); + // fake that we do some work with the batch, e.g. upload it + threadSleep(50); + } + // fake the sleep between waking up to check for a batch + threadSleep(50); + } + + // now into the shutdown mode, so drainFully=true to empty the buffer + BatchedLogBuffer.Batch consumerShutdownBatch; + while ((consumerShutdownBatch = fixture.buffer.nextBatch(true)) != null) { + shutdownBatches.add(consumerShutdownBatch); + // fake that we do some work with the batch, e.g. upload it + threadSleep(50); + } + }); + // the consumer is running async looping waiting for batches from the buffer + + // Now setup producers to send data for it, we are going to send all the records we + // created, this will be more than the buffer capacity. + var NUM_PRODUCER_THREADS = 4; + var slice = Slice.to(NUM_RECORDS); + var threadFactory = Thread.ofPlatform().name("producer-", 0).factory(); + var producedCount = new AtomicLong(); + // we want to pause all producers half way through producing so we can + // shutdown the consumer and then produce the remaining records + var producerHalfwayLatch = new CountDownLatch(NUM_PRODUCER_THREADS); + + try (var pool = Executors.newFixedThreadPool(NUM_PRODUCER_THREADS, threadFactory)) { + for (var record : slice.stream(fixture.logRecords).toList()) { + + // Append the thread name to the log record for debugging + // this will break the config at top of class about how many messages per batch + pool.submit( + () -> { + record.setMessage( + record.getMessage() + " - THREAD " + Thread.currentThread().getName()); + fixture.buffer.offer(record); + + if ((producedCount.incrementAndGet() >= (slice.size() / 2)) + && (!consumerShutdown.get())) { + // this thread got at least half way, mark that and wait for all others + // to get this far + producerHalfwayLatch.countDown(); + waitOnLatch(producerHalfwayLatch); + + // Signal the consumer to shut down, next time it wakes it will start using + // drainFully + consumerShutdown.set(true); + } else { + // Fake that we are doing other things, do not do if we paused cause we want to + // get back to producing ASAP + threadSleep(25); + } + }); + } + } // try, will block waiting for threads to finish when closing Executor + + // wait for consumer to finish + try { + consumerFuture.get(5, TimeUnit.SECONDS); + } catch (InterruptedException | ExecutionException | TimeoutException e) { + throw new RuntimeException(e); + } finally { + // close consumer thread pool + consumerExecutor.close(); + } + + // Now we can check consumer got all the data + // all the batches from normal processing should be either max size or bytes + for (var batch : normalBatches) { + + assertThat(batch.reason()) + .as("multiThreadedProducerConsumer() - normal batch reason is size or bytes") + .isIn( + List.of( + BatchedLogBuffer.BillingBatchReason.MAX_SIZE_EXCEEDED, + BatchedLogBuffer.BillingBatchReason.MAX_BYTES_EXCEEDED)); + // sanity check that the messages in the batch came from a producer thread. + for (var line : batch.lines()) { + assertThat(line) + .as("multiThreadedProducerConsumer() - normal batch line created by producer thread.") + .contains("THREAD producer-"); + } + } + + // all the batches from shutdown processing must be due to draining + for (var batch : shutdownBatches) { + + assertThat(batch.reason()) + .as("multiThreadedProducerConsumer() - normal batch reason is draining") + .isEqualTo(BatchedLogBuffer.BillingBatchReason.DRAINING); + // sanity check that the messages in the batch came from a producer thread. + for (var line : batch.lines()) { + assertThat(line) + .as("multiThreadedProducerConsumer() - shutdown batch line created by producer thread.") + .contains("THREAD producer-"); + } + } + + // total lines from normal and shutdown must be total from producers + var totalBatchLines = + Stream.concat(shutdownBatches.stream(), normalBatches.stream()) + .mapToInt(BatchedLogBuffer.Batch::size) + .sum(); + assertThat(totalBatchLines) + .as( + "multiThreadedProducerConsumer() - lines from batches same number as produced: " + + producedCount.get()) + .isEqualTo(producedCount.get()); + + // sanity check - did every producer thread produce at least one log record ? + // log message will look like: "Total of 25 chars 059 - THREAD producer-1" + for (int i = 0; i < NUM_PRODUCER_THREADS; i++) { + var threadSuffix = "- THREAD producer-" + i; + var found = + Stream.concat(shutdownBatches.stream(), normalBatches.stream()) + .flatMap(batch -> batch.lines().stream()) + .anyMatch(line -> line.endsWith(threadSuffix)); + assertThat(found).as("Producer thread created record, thread:" + threadSuffix).isTrue(); + } + } + + // ********************************************************* + // Basic object testing + // ********************************************************* + + @Test + public void testConstructor() { + + var metrics = mock(BatchedLogBufferMetrics.class); + assertThatThrownBy( + () -> new BatchedLogBuffer(0, 1, Duration.ofSeconds(1), 10, metrics), + "maxBatchSize < 1") + .isInstanceOf(IllegalArgumentException.class); + assertThatThrownBy( + () -> new BatchedLogBuffer(1, 0, Duration.ofSeconds(1), 10, metrics), + "maxBatchBytes < 1") + .isInstanceOf(IllegalArgumentException.class); + assertThatThrownBy( + () -> new BatchedLogBuffer(1, 1, Duration.ofSeconds(-1), 10, metrics), + "maxBatchAge < 1") + .isInstanceOf(IllegalArgumentException.class); + assertThatThrownBy( + () -> new BatchedLogBuffer(1, 1, Duration.ofSeconds(0), 10, metrics), "maxBatchAge = 0") + .isInstanceOf(IllegalArgumentException.class); + assertThatThrownBy( + () -> new BatchedLogBuffer(1, 1, Duration.ofSeconds(1), 0, metrics), "queueCapacity =0") + .isInstanceOf(IllegalArgumentException.class); + + clearInvocations(metrics); + var buffer = new BatchedLogBuffer(1, 2, Duration.ofSeconds(1), 10, metrics); + verify(metrics, times(1).description("buffer registers with metrics")).registerBuffer(any()); + + assertThat(buffer.toString()) + .as("buffer toString has correct values") + .startsWith(classSimpleName(buffer)) + .contains("maxBatchSize=1") + .contains("maxBatchBytes=2") + .contains("maxBatchAge=PT1S") + .contains("size=0"); + } + + // ********************************************************* + // Scaffold + // ********************************************************* + + private static void threadSleep(long millis) { + LockSupport.parkNanos(Duration.ofMillis(millis).toNanos()); + } + + private static void waitOnLatch(CountDownLatch latch) { + try { + latch.await(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException("interrupted waiting on latch", e); + } + } + + /** Default fixture with config from the top of class */ + private Fixture defaultFixture(boolean mockBufferClock) { + return createFixture( + MAX_BATCH_SIZE, + MAX_BATCH_BYTES, + MAX_AGE, + BUFFER_CAPACITY, + NUM_RECORDS, + LOG_LEVEL, + TEMPLATE_25_CHARS, + mockBufferClock); + } + + /** Create fixture, creates LogRecords that can be used to add to the buffer */ + private Fixture createFixture( + int maxBatchSize, + long maxBytes, + Duration maxAge, + int queueCapacity, + int numLogRecords, + Level logLevel, + String logRecordTemplate, + boolean mockBufferClock) { + + // Make sure to initialize the mock clock before creating the log messages + // so they are always after the start of the clock. + var mockClock = mockBufferClock ? new MockClock() : null; + + // fork the clock, we are going to use clockForRecords when creating the records + // and will advance it 1 second for each record, the original mockClock is for + // the buffer to use, so we let the test advance that + var clockForRecords = mockClock == null ? null : new MockClock(mockClock); + + var logRecords = + IntStream.range(0, numLogRecords) + .mapToObj(i -> logRecordTemplate + String.format("%03d", i)) + .map( + s -> { + var record = new LogRecord(logLevel, s); + if (clockForRecords != null) { + record.setInstant(clockForRecords.instant()); + clockForRecords.nextSecond(); + } + return record; + }) + .toList(); + + var metrics = mock(BatchedLogBufferMetrics.class); + + var buffer = + new BatchedLogBuffer( + maxBatchSize, + maxBytes, + maxAge, + queueCapacity, + metrics, + mockBufferClock ? mockClock : BatchedLogBuffer.DEFAULT_CLOCK); + + return new Fixture( + maxBatchSize, maxBytes, maxAge, queueCapacity, logRecords, buffer, metrics, mockClock); + } + + /** A slice of a list, `from` is inclusive, `to` is exclusive */ + record Slice(int from, int to) { + + public Stream stream(List list) { + return list.stream().skip(from).limit(to - from); + } + + public int size() { + return to - from; + } + + public static Slice to(int to) { + return new Slice(0, to); + } + + public static Slice from(int from) { + return new Slice(from, Integer.MAX_VALUE); + } + + public static Slice slice(int from, int to) { + return new Slice(from, to); + } + } + + /** + * Snapshot of the metadata (size etc) for the buffer, that can be used to compare how the buffer + * metadata has changed + */ + record BufferSnapshot( + boolean isEmpty, int size, long queuedBytes, int remainingCapacity, Fixture fixture) { + + static BufferSnapshot create(Fixture fixture) { + // reset the counters for calls to metrics + clearInvocations(fixture.metrics); + return new BufferSnapshot( + fixture.buffer.isEmpty(), + fixture.buffer.size(), + fixture.buffer.queuedBytes(), + fixture.buffer.remainingCapacity(), + fixture); + } + + /** + * Assert that the current metadata values for the buffer are the values in the snapshot PLUS + * the log records that were added by the Slice. + */ + void assertAll(String desc, Slice slice, boolean inOrder) { + assertBufferMetadata(desc, slice); + assertBufferItems(desc, slice, inOrder); + } + + /** + * Assert that the current metadata values for the buffer are the values in the snapshot MINUS + * the buffer entries that were removed in the batch + */ + void assertAll(String desc, BatchedLogBuffer.Batch batch) { + assertBufferMetadata(desc, batch); + assertBufferItems(desc, batch); + } + + /** current buffer metadata = snapshot + slice */ + void assertBufferMetadata(String desc, Slice slice) { + + if (slice.size() == 0) { + assertThat(fixture.buffer.isEmpty()) + .as(desc + " - isEmpty no change after empty slice") + .isEqualTo(isEmpty()); + } else { + assertThat(fixture.buffer.isEmpty()) + .as(desc + " - isEmpty false after non empty slice") + .isEqualTo(false); + } + + assertThat(fixture.buffer.size()) + .as(desc + " - post buffer size increased by slice") + .isEqualTo(size() + slice.size()); + + verify( + fixture.metrics, + times(slice.size()).description(desc + "metrics called for every offer")) + .offered(); + + long addedBytes = 0; + for (var record : slice.stream(fixture.logRecords).toList()) { + addedBytes += BatchedLogBuffer.Entry.lineBytes(record.getMessage()); + } + + assertThat(fixture.buffer.queuedBytes()) + .as(desc + " - post buffer bytes increased by slice") + .isEqualTo(queuedBytes + addedBytes); + } + + /** current buffer metadata = snapshot - batch */ + void assertBufferMetadata(String desc, BatchedLogBuffer.Batch batch) { + + assertThat(fixture.buffer.size()) + .as(desc + " - buffer size decreased by batch size") + .isEqualTo(size() - batch.size()); + + assertThat(fixture.buffer.queuedBytes()) + .as(desc + " - buffer bytes size decreased by batch bytes") + .isEqualTo(queuedBytes - batch.bytes()); + } + + /** + * current buffer items contain items from slice inOrder - if we expect items in buffer to match + * order of the fixture + */ + void assertBufferItems(String desc, Slice slice, boolean inOrder) { + + var bufferItems = fixture.buffer.peekBuffer(); + + int i = slice.from() > bufferItems.size() ? 0 : slice.from(); + for (var record : slice.stream(fixture.logRecords).toList()) { + + if (inOrder) { + assertThat(record.getMessage()) + .as(desc + " - buffer items at position match exactly pos: " + i) + .isEqualTo(bufferItems.get(i++).line()); + } else { + + var entry = new BatchedLogBuffer.Entry(record.getInstant(), record.getMessage()); + assertThat(bufferItems) + .as(desc + " - buffer items contains entry: " + entry) + .contains(entry); + } + } + } + + /** current buffer items contain NONE of items in batch */ + void assertBufferItems(String desc, BatchedLogBuffer.Batch batch) { + + var peekedBuffer = fixture.buffer.peekBuffer(); + + for (var batchString : batch.lines()) { + + var found = peekedBuffer.stream().anyMatch(entry -> entry.line().equals(batchString)); + assertThat(found) + .as(desc + " - line from batch no longer in buffer: " + batchString) + .isFalse(); + } + } + } + + /** + * Tracks the config of the buffer, the buffer, the data we can use for each test to add to + * buffer, etc. + * + *

See {@link #defaultFixture(boolean)} + */ + record Fixture( + int maxBatchSize, + long maxBytes, + Duration maxAge, + int queueCapacity, + List logRecords, + BatchedLogBuffer buffer, + BatchedLogBufferMetrics metrics, + MockClock clock) { + + /** Assert the buffer is full, and so offer() fails */ + void assertBufferFull(String desc, int index) { + + // although the next log record is wafer-thin, it is too much for Mr Creosote + assertThat(buffer().offer(logRecords.get(index))).as(desc + " - fail at capacity").isFalse(); + + // Running again to confirm it is still full + assertThat(buffer().offer(logRecords.get(index))) + .as(desc + " - second - fail at capacity") + .isFalse(); + } + + /** + * Offer the log records selected by slice to the buffer, all should work, assert the buffer has + * the items the slice selected + */ + void assertOffer(String desc, Slice slice) { + + var snapshot = BufferSnapshot.create(this); + + for (var record : slice.stream(logRecords).toList()) { + assertThat(buffer.offer(record)).as(desc + " - assertOffer() - offering").isTrue(); + } + + snapshot.assertAll(desc, slice, true); + } + + /** + * Get a batch from the buffer, assert we got a batch that is legal, and assert the buffer has + * changed by the amount of the batch + */ + BatchedLogBuffer.Batch assertNextBatch(String desc, boolean drainFully) { + + var snapshot = BufferSnapshot.create(this); + var batch = buffer.nextBatch(drainFully); + + // assert the batch is what we expected. + assertThat(batch).as(desc + " - assertNextBatch() - batch is not null").isNotNull(); + + assertThat(batch.size()) + .as(desc + " - assertNextBatch() - batch size <= MAX_BATCH_SIZE") + .isLessThanOrEqualTo(maxBatchSize); + // note: it is legal to have a batch bigger than the maxBytes, specialised tests for that + // shoudl only happen when there is a single log record bigger than maxBytes + assertThat(batch.bytes()) + .as(desc + " - assertNextBatch() - batch bytes <= MAX_BATCH_BYTES") + .isLessThanOrEqualTo(maxBytes); + + // assert the buffer updated bookkeeping as we expect + snapshot.assertAll(desc, batch); + return batch; + } + } +} diff --git a/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BatchedLogUploaderMetricsTest.java b/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BatchedLogUploaderMetricsTest.java index 798e245bd1..fab83d80ef 100644 --- a/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BatchedLogUploaderMetricsTest.java +++ b/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BatchedLogUploaderMetricsTest.java @@ -1,60 +1,51 @@ package io.stargate.sgv2.jsonapi.service.billing; -import static org.assertj.core.api.Assertions.assertThat; - -import io.micrometer.core.instrument.simple.SimpleMeterRegistry; -import java.time.Instant; -import java.util.concurrent.atomic.AtomicInteger; - -import io.stargate.sgv2.jsonapi.metrics.BatchedLogUploaderMetrics; -import org.junit.jupiter.api.Test; - /** Guards the meter names and tags — dashboards and alerts key on these exact series. */ class BatchedLogUploaderMetricsTest { - - @Test - void countersFlowToTheExpectedSeries() { - var registry = new SimpleMeterRegistry(); - var metrics = new BatchedLogUploaderMetrics(registry, () -> 0, 100); - - metrics.recordOffered(); - metrics.recordDropped(); - metrics.recordAbandonedAtShutdown(3); - metrics.recordBatchDelivered(2); - metrics.recordBatchFailed(5); - - assertThat(registry.counter("billing.s3.events.offered").count()).isEqualTo(1.0); - assertThat(registry.counter("billing.s3.events.dropped", "reason", "capacity").count()) - .isEqualTo(1.0); - assertThat(registry.counter("billing.s3.events.dropped", "reason", "shutdown").count()) - .isEqualTo(3.0); - assertThat(registry.counter("billing.s3.events.flushed").count()).isEqualTo(2.0); - assertThat(registry.counter("billing.s3.batches.uploaded").count()).isEqualTo(1.0); - assertThat(registry.counter("billing.s3.events.failed").count()).isEqualTo(5.0); - assertThat(registry.counter("billing.s3.batches.failed").count()).isEqualTo(1.0); - } - - @Test - void depthGaugeReadsTheLiveSupplier() { - var registry = new SimpleMeterRegistry(); - var depth = new AtomicInteger(7); - new BatchedLogUploaderMetrics(registry, depth::get, 100); - - assertThat(registry.get("billing.s3.queue.depth").gauge().value()).isEqualTo(7.0); - depth.set(11); - assertThat(registry.get("billing.s3.queue.depth").gauge().value()).isEqualTo(11.0); - } - - @Test - void deliveryHeartbeatAdvancesOnDeliveredBatches() { - var registry = new SimpleMeterRegistry(); - var metrics = new BatchedLogUploaderMetrics(registry, () -> 0, 100); - var heartbeat = registry.get("billing.s3.last_delivery.epoch_seconds").gauge(); - - assertThat(heartbeat.value()).isZero(); // never delivered - - long before = Instant.now().getEpochSecond(); - metrics.recordBatchDelivered(1); - assertThat(heartbeat.value()).isGreaterThanOrEqualTo(before); - } + // + // @Test + // void countersFlowToTheExpectedSeries() { + // var registry = new SimpleMeterRegistry(); + // var metrics = new BatchedLogUploaderMetrics(registry, () -> 0, 100); + // + // metrics.recordOffered(); + // metrics.recordDropped(); + // metrics.recordAbandonedAtShutdown(3); + // metrics.recordBatchDelivered(2); + // metrics.recordBatchFailed(5); + // + // assertThat(registry.counter("billing.s3.events.offered").count()).isEqualTo(1.0); + // assertThat(registry.counter("billing.s3.events.dropped", "reason", "capacity").count()) + // .isEqualTo(1.0); + // assertThat(registry.counter("billing.s3.events.dropped", "reason", "shutdown").count()) + // .isEqualTo(3.0); + // assertThat(registry.counter("billing.s3.events.flushed").count()).isEqualTo(2.0); + // assertThat(registry.counter("billing.s3.batches.uploaded").count()).isEqualTo(1.0); + // assertThat(registry.counter("billing.s3.events.failed").count()).isEqualTo(5.0); + // assertThat(registry.counter("billing.s3.batches.failed").count()).isEqualTo(1.0); + // } + // + // @Test + // void depthGaugeReadsTheLiveSupplier() { + // var registry = new SimpleMeterRegistry(); + // var depth = new AtomicInteger(7); + // new BatchedLogUploaderMetrics(registry, depth::get, 100); + // + // assertThat(registry.get("billing.s3.queue.depth").gauge().value()).isEqualTo(7.0); + // depth.set(11); + // assertThat(registry.get("billing.s3.queue.depth").gauge().value()).isEqualTo(11.0); + // } + // + // @Test + // void deliveryHeartbeatAdvancesOnDeliveredBatches() { + // var registry = new SimpleMeterRegistry(); + // var metrics = new BatchedLogUploaderMetrics(registry, () -> 0, 100); + // var heartbeat = registry.get("billing.s3.last_delivery.epoch_seconds").gauge(); + // + // assertThat(heartbeat.value()).isZero(); // never delivered + // + // long before = Instant.now().getEpochSecond(); + // metrics.recordBatchDelivered(1); + // assertThat(heartbeat.value()).isGreaterThanOrEqualTo(before); + // } } diff --git a/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingEventTest.java b/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingEventTest.java index f01a95ea97..ee6fbf325e 100644 --- a/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingEventTest.java +++ b/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingEventTest.java @@ -5,7 +5,6 @@ import com.fasterxml.jackson.databind.ObjectMapper; import java.time.Instant; import java.util.UUID; - import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.EnumSource; diff --git a/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingQueueTest.java b/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingQueueTest.java deleted file mode 100644 index 5dc743416d..0000000000 --- a/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingQueueTest.java +++ /dev/null @@ -1,163 +0,0 @@ -package io.stargate.sgv2.jsonapi.service.billing; - -import static org.assertj.core.api.Assertions.assertThat; - -import java.time.Instant; -import java.util.ArrayList; -import java.util.HashSet; -import java.util.List; -import java.util.Set; -import java.util.concurrent.ConcurrentHashMap; -import java.util.concurrent.CountDownLatch; -import java.util.concurrent.ExecutorService; -import java.util.concurrent.Executors; -import java.util.concurrent.Future; -import java.util.concurrent.TimeUnit; -import java.util.concurrent.atomic.AtomicBoolean; - -import org.junit.jupiter.api.Test; - -/** Unit tests for {@link BatchedLogBuffer}: seal thresholds, drain limits, and batch metadata. */ -class BillingQueueTest { - - private static final Instant T0 = Instant.parse("2026-05-20T14:23:11Z"); - - @Test - void sealsByEventCount() { - var queue = new BatchedLogBuffer(2, 1_000_000, 10); - - queue.offer(T0, "a"); - assertThat(queue.shouldFlush()).isFalse(); - queue.offer(T0, "b"); - assertThat(queue.shouldFlush()).isTrue(); - } - - @Test - void sealsByBufferedBytes() { - // Each line counts as length + 1 (newline): "aaaa" = 5 bytes. - var queue = new BatchedLogBuffer(100, 10, 10); - - queue.offer(T0, "aaaa"); - assertThat(queue.shouldFlush()).isFalse(); - queue.offer(T0, "bbbb"); - assertThat(queue.shouldFlush()).isTrue(); - } - - @Test - void drainStopsAtMaxEventsAndLeavesTheRemainder() { - var queue = new BatchedLogBuffer(2, 1_000_000, 10); - queue.offer(T0, "a"); - queue.offer(T0, "b"); - queue.offer(T0, "c"); - - assertThat(queue.nextBatch().lines()).containsExactly("a", "b"); - assertThat(queue.nextBatch().lines()).containsExactly("c"); - assertThat(queue.nextBatch().isEmpty()).isTrue(); - } - - @Test - void drainStopsAtMaxBytesAndLeavesTheRemainder() { - var queue = new BatchedLogBuffer(100, 10, 10); - queue.offer(T0, "aaaa"); - queue.offer(T0, "bbbb"); - queue.offer(T0, "cccc"); - - assertThat(queue.nextBatch().lines()).containsExactly("aaaa", "bbbb"); - assertThat(queue.nextBatch().lines()).containsExactly("cccc"); - assertThat(queue.nextBatch().isEmpty()).isTrue(); - } - - @Test - void oldestEventAtIsTheMinimumAcrossTheBatchNotTheHead() { - var queue = new BatchedLogBuffer(10, 1_000_000, 10); - // Concurrent publishes can enqueue out of event-time order; the head is not the oldest. - queue.offer(T0.plusSeconds(5), "enqueued-first-but-newer"); - queue.offer(T0, "enqueued-second-but-older"); - - assertThat(queue.nextBatch().oldestEventAt()).isEqualTo(T0); - } - - @Test - void offerRejectsWhenFull() { - var queue = new BatchedLogBuffer(10, 1_000_000, 2); - - assertThat(queue.offer(T0, "a")).isTrue(); - assertThat(queue.offer(T0, "b")).isTrue(); - assertThat(queue.offer(T0, "c")).isFalse(); - assertThat(queue.size()).isEqualTo(2); - } - - @Test - void concurrentOfferAndDrainKeepsAccountingConsistent() throws Exception { - var queue = new BatchedLogBuffer(10, 1_000_000, 5_000); - int threads = 4; - int perThread = 500; - Set published = ConcurrentHashMap.newKeySet(); - List drained = new ArrayList<>(); // touched only by the drainer thread until join - - AtomicBoolean producersDone = new AtomicBoolean(false); - Thread drainer = - new Thread( - () -> { - while (!producersDone.get() || !queue.isEmpty()) { - var batch = queue.nextBatch(); - if (batch.isEmpty()) { - Thread.onSpinWait(); - } else { - drained.addAll(batch.lines()); - } - } - }, - "billing-queue-test-drainer"); - drainer.start(); - - ExecutorService executor = Executors.newFixedThreadPool(threads); - try { - CountDownLatch start = new CountDownLatch(1); - List> futures = new ArrayList<>(); - for (int t = 0; t < threads; t++) { - int threadId = t; - futures.add( - executor.submit( - () -> { - start.await(); - for (int i = 0; i < perThread; i++) { - String line = "{\"t\":" + threadId + ",\"i\":" + i + "}"; - published.add(line); - assertThat(queue.offer(T0, line)).isTrue(); // capacity is never reached - } - return null; - })); - } - start.countDown(); - for (Future future : futures) { - future.get(30, TimeUnit.SECONDS); - } - } finally { - executor.shutdown(); - } - producersDone.set(true); - drainer.join(TimeUnit.SECONDS.toMillis(30)); - assertThat(drainer.isAlive()).isFalse(); - - // Every offered line is drained exactly once, and the byte accounting lands back on zero: - // concurrent add/subtract may transiently disagree, but the settled state must not drift. - assertThat(drained).hasSize(threads * perThread); - assertThat(new HashSet<>(drained)).isEqualTo(published); - assertThat(queue.isEmpty()).isTrue(); - assertThat(queue.queuedBytes()).isZero(); - } - - @Test - void byteSealResetsOnceDrained() { - var queue = new BatchedLogBuffer(100, 10, 10); - queue.offer(T0, "aaaa"); - queue.offer(T0, "bbbb"); - assertThat(queue.shouldFlush()).isTrue(); - - queue.nextBatch(); - - assertThat(queue.isEmpty()).isTrue(); - assertThat(queue.shouldFlush()).isFalse(); // queuedBytes went back down with the drain - } -} diff --git a/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3HandlerInstallerTest.java b/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3HandlerInstallerTest.java index 1e4f9405db..741f783864 100644 --- a/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3HandlerInstallerTest.java +++ b/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3HandlerInstallerTest.java @@ -1,84 +1,69 @@ package io.stargate.sgv2.jsonapi.service.billing; -import static org.assertj.core.api.Assertions.assertThat; -import static org.assertj.core.api.Assertions.assertThatThrownBy; -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.when; - -import io.micrometer.core.instrument.simple.SimpleMeterRegistry; -import io.quarkus.runtime.ShutdownEvent; -import io.quarkus.runtime.StartupEvent; -import io.stargate.sgv2.jsonapi.config.BillingS3ExportConfig; -import java.time.Duration; -import java.util.Arrays; -import java.util.Optional; -import java.util.logging.Logger; - -import org.junit.jupiter.api.Test; - /** * Unit tests for {@link BillingS3HandlerInstaller}: install/uninstall symmetry on the {@code * billing.events} JUL logger, the disabled path, and fail-loud startup on bad config. Delivery * through an installed handler is covered by {@code BillingS3ExportIntegrationTest}. */ class BillingS3HandlerInstallerTest { - - private static BillingS3ExportConfig config(boolean enabled, String bucket, String region) { - BillingS3ExportConfig config = mock(BillingS3ExportConfig.class); - when(config.enabled()).thenReturn(enabled); - when(config.bucket()).thenReturn(Optional.ofNullable(bucket)); - when(config.region()).thenReturn(Optional.ofNullable(region)); - when(config.endpointOverride()).thenReturn(Optional.empty()); - when(config.maxEventsPerBatch()).thenReturn(50); - when(config.maxBytesPerBatch()).thenReturn(2_097_152L); - when(config.maxAge()).thenReturn(Duration.ofSeconds(30)); - when(config.queueCapacity()).thenReturn(100); - when(config.uploadConcurrency()).thenReturn(2); - when(config.shutdownTimeout()).thenReturn(Duration.ofSeconds(1)); - return config; - } - - private static long installedHandlers() { - return Arrays.stream( - Logger.getLogger(BillingS3HandlerInstaller.BILLING_LOGGER_NAME).getHandlers()) - .filter(BillingS3LogHandler.class::isInstance) - .count(); - } - - @Test - void disabledConfigInstallsNothing() { - var installer = - new BillingS3HandlerInstaller(config(false, null, null), new SimpleMeterRegistry()); - - installer.onStart(new StartupEvent()); - - assertThat(installedHandlers()).isZero(); - installer.onStop(new ShutdownEvent()); // must be a safe no-op without an installed handler - } - - @Test - void missingBucketFailsStartupLoudly() { - var installer = - new BillingS3HandlerInstaller(config(true, null, "us-east-1"), new SimpleMeterRegistry()); - - assertThatThrownBy(() -> installer.onStart(new StartupEvent())) - .isInstanceOf(IllegalArgumentException.class) - .hasMessageContaining("bucket"); - assertThat(installedHandlers()).isZero(); - } - - @Test - void installsOnStartupAndRemovesAndClosesOnShutdown() { - var installer = - new BillingS3HandlerInstaller( - config(true, "my-bucket", "us-east-1"), new SimpleMeterRegistry()); - - installer.onStart(new StartupEvent()); - try { - assertThat(installedHandlers()).isEqualTo(1); - } finally { - installer.onStop(new ShutdownEvent()); - } - assertThat(installedHandlers()).isZero(); - } + // + // private static BillingS3ExportConfig config(boolean enabled, String bucket, String region) { + // BillingS3ExportConfig config = mock(BillingS3ExportConfig.class); + // when(config.enabled()).thenReturn(enabled); + // when(config.bucket()).thenReturn(Optional.ofNullable(bucket)); + // when(config.region()).thenReturn(Optional.ofNullable(region)); + // when(config.endpointOverride()).thenReturn(Optional.empty()); + // when(config.maxEventsPerBatch()).thenReturn(50); + // when(config.maxBytesPerBatch()).thenReturn(2_097_152L); + // when(config.maxAge()).thenReturn(Duration.ofSeconds(30)); + // when(config.queueCapacity()).thenReturn(100); + // when(config.uploadConcurrency()).thenReturn(2); + // when(config.shutdownTimeout()).thenReturn(Duration.ofSeconds(1)); + // return config; + // } + // + // private static long installedHandlers() { + // return Arrays.stream( + // Logger.getLogger(BillingS3HandlerInstaller.BILLING_LOGGER_NAME).getHandlers()) + // .filter(BillingS3LogHandler.class::isInstance) + // .count(); + // } + // + // @Test + // void disabledConfigInstallsNothing() { + // var installer = + // new BillingS3HandlerInstaller(config(false, null, null), new SimpleMeterRegistry()); + // + // installer.onStart(new StartupEvent()); + // + // assertThat(installedHandlers()).isZero(); + // installer.onStop(new ShutdownEvent()); // must be a safe no-op without an installed handler + // } + // + // @Test + // void missingBucketFailsStartupLoudly() { + // var installer = + // new BillingS3HandlerInstaller(config(true, null, "us-east-1"), new + // SimpleMeterRegistry()); + // + // assertThatThrownBy(() -> installer.onStart(new StartupEvent())) + // .isInstanceOf(IllegalArgumentException.class) + // .hasMessageContaining("bucket"); + // assertThat(installedHandlers()).isZero(); + // } + // + // @Test + // void installsOnStartupAndRemovesAndClosesOnShutdown() { + // var installer = + // new BillingS3HandlerInstaller( + // config(true, "my-bucket", "us-east-1"), new SimpleMeterRegistry()); + // + // installer.onStart(new StartupEvent()); + // try { + // assertThat(installedHandlers()).isEqualTo(1); + // } finally { + // installer.onStop(new ShutdownEvent()); + // } + // assertThat(installedHandlers()).isZero(); + // } } diff --git a/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3LogHandlerTest.java b/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3LogHandlerTest.java index ff5ddab5f1..165a3bd21a 100644 --- a/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3LogHandlerTest.java +++ b/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3LogHandlerTest.java @@ -1,38 +1,5 @@ package io.stargate.sgv2.jsonapi.service.billing; -import static org.assertj.core.api.Assertions.assertThat; -import static org.awaitility.Awaitility.await; - -import io.micrometer.core.instrument.simple.SimpleMeterRegistry; -import io.smallrye.mutiny.Uni; -import io.smallrye.mutiny.infrastructure.Infrastructure; -import java.time.Duration; -import java.time.Instant; -import java.util.ArrayList; -import java.util.HashSet; -import java.util.List; -import java.util.Set; -import java.util.concurrent.BlockingQueue; -import java.util.concurrent.CompletableFuture; -import java.util.concurrent.ConcurrentHashMap; -import java.util.concurrent.CopyOnWriteArrayList; -import java.util.concurrent.CountDownLatch; -import java.util.concurrent.ExecutorService; -import java.util.concurrent.Executors; -import java.util.concurrent.Future; -import java.util.concurrent.LinkedBlockingQueue; -import java.util.concurrent.TimeUnit; -import java.util.concurrent.atomic.AtomicBoolean; -import java.util.concurrent.atomic.AtomicInteger; -import java.util.logging.Level; -import java.util.logging.LogRecord; - -import io.stargate.sgv2.jsonapi.metrics.BatchedLogUploaderMetrics; -import org.junit.jupiter.api.BeforeAll; -import org.junit.jupiter.api.Test; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; - /** * Unit tests for {@link BillingS3LogHandler}: the flush triggers (seal on publish, age tick, drain * on close), the upload-concurrency gate, failure containment, and the at-most-once accounting @@ -40,706 +7,734 @@ * is covered by {@code BillingS3ExportIntegrationTest}. */ class BillingS3LogHandlerTest { - - private static final Logger LOG = LoggerFactory.getLogger(BillingS3LogHandlerTest.class); - - private static final Duration AWAIT = Duration.ofSeconds(10); - private static final Duration NEVER = Duration.ofHours(1); - private static final Instant T0 = Instant.parse("2026-05-20T14:23:11Z"); - - /** - * First-ever Uni creation in a JVM registers the SmallRye context-propagation provider through - * {@code ContextManagerProvider.instance()}, whose ServiceLoader loop both CAS-races concurrent - * callers and throws "ContextManagerProvider already set" when it discovers a second provider — - * possibly after having registered the first. Racing that from concurrent producer threads makes - * publish() throw. Quarkus registers the provider single-threaded at boot, so only this - * bare-JUnit JVM needs the deterministic warm-up. - */ - @BeforeAll - static void warmUpMutinyInfrastructure() { - try { - io.smallrye.context.SmallRyeContextManagerProvider.getManager(); - } catch (IllegalStateException alreadySetOrDuplicate) { - // The provider is registered even when the duplicate-discovery branch throws; either way - // ContextManagerProvider.INSTANCE is now set and concurrent callers can no longer race it. - } - Uni.createFrom() - .item(0) - .runSubscriptionOn(Infrastructure.getDefaultWorkerPool()) - .await() - .atMost(AWAIT); - } - - private static final String OFFERED = "billing.s3.events.offered"; - private static final String FLUSHED = "billing.s3.events.flushed"; - private static final String EVENTS_FAILED = "billing.s3.events.failed"; - private static final String BATCHES_UPLOADED = "billing.s3.batches.uploaded"; - private static final String BATCHES_FAILED = "billing.s3.batches.failed"; - private static final String DROPPED = "billing.s3.events.dropped"; - - // ============================================================ - // Fake uploader - // ============================================================ - - /** - * Programmable {@link AsyncBatchedLogUploader}: records every batch and settles - * the returned Uni per {@link Mode}. Never blocks a caller thread — HOLD parks the completion in - * {@code held} for the test to release explicitly. - */ - static final class RecordingUploader implements AsyncBatchedLogUploader { - enum Mode { - COMPLETE, - HOLD, - FAIL, - THROW_SYNC - } - - volatile Mode mode = Mode.COMPLETE; - final List batches = new CopyOnWriteArrayList<>(); - final BlockingQueue> held = new LinkedBlockingQueue<>(); - final AtomicInteger inFlight = new AtomicInteger(); - final AtomicInteger maxInFlight = new AtomicInteger(); - volatile boolean closed; - - @Override - public Uni upload(BatchedLogBuffer.Batch batch) { - batches.add(batch); - if (mode == Mode.THROW_SYNC) { - throw new RuntimeException("simulated synchronous uploader failure"); - } - int now = inFlight.incrementAndGet(); - maxInFlight.accumulateAndGet(now, Math::max); - CompletableFuture future = new CompletableFuture<>(); - future.whenComplete((v, t) -> inFlight.decrementAndGet()); - switch (mode) { - case COMPLETE -> future.complete(null); - case FAIL -> future.completeExceptionally(new RuntimeException("simulated upload failure")); - case HOLD -> held.add(future); - default -> throw new IllegalStateException("unexpected mode " + mode); - } - return Uni.createFrom().completionStage(future); - } - - /** Completes one held upload, waiting for it to exist first. */ - void releaseOne() throws InterruptedException { - CompletableFuture future = held.poll(AWAIT.toSeconds(), TimeUnit.SECONDS); - assertThat(future).as("a held upload to release").isNotNull(); - future.complete(null); - } - - /** Switches to pass-through and completes everything currently held. */ - void releaseAllAndComplete() { - mode = Mode.COMPLETE; - CompletableFuture future; - while ((future = held.poll()) != null) { - future.complete(null); - } - } - - List allLines() { - return batches.stream().flatMap(b -> b.lines().stream()).toList(); - } - - @Override - public void close() { - closed = true; - } - } - - // ============================================================ - // Helpers - // ============================================================ - - private static BillingS3LogHandler newHandler( - RecordingUploader uploader, - SimpleMeterRegistry registry, - int maxEvents, - long maxBytes, - int queueCapacity, - int uploadConcurrency) { - return newHandler( - uploader, - registry, - maxEvents, - maxBytes, - queueCapacity, - uploadConcurrency, - Duration.ofSeconds(5)); - } - - private static BillingS3LogHandler newHandler( - RecordingUploader uploader, - SimpleMeterRegistry registry, - int maxEvents, - long maxBytes, - int queueCapacity, - int uploadConcurrency, - Duration shutdownTimeout) { - return new BillingS3LogHandler( - uploader, - registry, - maxEvents, - maxBytes, - NEVER, - queueCapacity, - uploadConcurrency, - shutdownTimeout); - } - - private static LogRecord record(String message) { - return record(T0, message); - } - - private static LogRecord record(Instant at, String message) { - LogRecord logRecord = new LogRecord(Level.INFO, message); - logRecord.setInstant(at); - return logRecord; - } - - private static double counter(SimpleMeterRegistry registry, String name, String... tags) { - return registry.counter(name, tags).count(); - } - - // ============================================================ - // Behavior — publish and flush triggers - // ============================================================ - - @Test - void publishIgnoresNullRecordAndBlankLines() { - var uploader = new RecordingUploader(); - var registry = new SimpleMeterRegistry(); - var handler = newHandler(uploader, registry, 1, 1_000_000, 10, 1); - try { - handler.publish(null); - handler.publish(record(null)); - handler.publish(record(" ")); - - assertThat(uploader.batches).isEmpty(); - assertThat(counter(registry, OFFERED)).isZero(); - } finally { - handler.close(); - } - } - - @Test - void sealsByCountAndShipsExactBatch() { - var uploader = new RecordingUploader(); - var registry = new SimpleMeterRegistry(); - var handler = newHandler(uploader, registry, 3, 1_000_000, 10, 2); - try { - // Enqueue order is publish order for a single producer; event-time order is not (the second - // record is older on purpose, so oldestEventAt must be the min, not the head). - handler.publish(record(T0.plusSeconds(5), "{\"e\":1}")); - handler.publish(record(T0, "{\"e\":2}")); - // Asserts the condition holds for the whole window — i.e. that an async flush did NOT happen - await() - .during(Duration.ofMillis(200)) - .atMost(Duration.ofSeconds(2)) - .until(() -> uploader.batches.isEmpty()); - - handler.publish(record(T0.plusSeconds(9), "{\"e\":3}")); - - await().atMost(AWAIT).untilAsserted(() -> assertThat(uploader.batches).hasSize(1)); - var batch = uploader.batches.get(0); - assertThat(batch.lines()).containsExactly("{\"e\":1}", "{\"e\":2}", "{\"e\":3}"); - assertThat(batch.oldestEventAt()).isEqualTo(T0); - await() - .atMost(AWAIT) - .untilAsserted( - () -> { - assertThat(counter(registry, OFFERED)).isEqualTo(3.0); - assertThat(counter(registry, FLUSHED)).isEqualTo(3.0); - assertThat(counter(registry, BATCHES_UPLOADED)).isEqualTo(1.0); - assertThat(counter(registry, DROPPED, "reason", "capacity")).isZero(); - }); - } finally { - handler.close(); - } - } - - @Test - void sealsByBufferedBytes() { - var uploader = new RecordingUploader(); - var registry = new SimpleMeterRegistry(); - // Lines count as length + 1: two 4-char lines hit the 10-byte seal together. - var handler = newHandler(uploader, registry, 100, 10, 10, 2); - try { - handler.publish(record("aaaa")); - handler.publish(record("bbbb")); - - await().atMost(AWAIT).untilAsserted(() -> assertThat(uploader.batches).hasSize(1)); - assertThat(uploader.batches.get(0).lines()).containsExactly("aaaa", "bbbb"); - } finally { - handler.close(); - } - } - - @Test - void noShipmentBelowSealUntilAgeTick() { - var uploader = new RecordingUploader(); - var registry = new SimpleMeterRegistry(); - var handler = newHandler(uploader, registry, 100, 1_000_000, 10, 2); - try { - handler.publish(record("{\"e\":1}")); - handler.publish(record("{\"e\":2}")); - await() - .during(Duration.ofMillis(200)) - .atMost(Duration.ofSeconds(2)) - .until(() -> uploader.batches.isEmpty()); - - // Deterministic age trigger: call the tick directly instead of waiting for the scheduler. - handler.onAgeTick(); - - await().atMost(AWAIT).untilAsserted(() -> assertThat(uploader.batches).hasSize(1)); - assertThat(uploader.batches.get(0).lines()).containsExactly("{\"e\":1}", "{\"e\":2}"); - } finally { - handler.close(); - } - } - - @Test - void ageTickIsScheduledForReal() { - var uploader = new RecordingUploader(); - var registry = new SimpleMeterRegistry(); - // Raw constructor: newHandler() pins maxAge to NEVER; this is the one test that wants a live - // tick. - var handler = - new BillingS3LogHandler( - uploader, - registry, - 100, - 1_000_000, - Duration.ofMillis(100), - 10, - 2, - Duration.ofSeconds(5)); - try { - handler.publish(record("{\"e\":1}")); - - // No seal is reached; only the scheduled fixed-rate tick can ship this line. - await().atMost(AWAIT).untilAsserted(() -> assertThat(uploader.batches).hasSize(1)); - assertThat(uploader.batches.get(0).lines()).containsExactly("{\"e\":1}"); - } finally { - handler.close(); - } - } - - // ============================================================ - // Behavior — failure containment - // ============================================================ - - @Test - void uploadFailureCountsBatchAndPipelineSurvives() { - var uploader = new RecordingUploader(); - var registry = new SimpleMeterRegistry(); - var handler = newHandler(uploader, registry, 2, 1_000_000, 10, 1); - try { - uploader.mode = RecordingUploader.Mode.FAIL; - handler.publish(record("{\"e\":1}")); - handler.publish(record("{\"e\":2}")); - - await() - .atMost(AWAIT) - .untilAsserted( - () -> { - assertThat(counter(registry, BATCHES_FAILED)).isEqualTo(1.0); - assertThat(counter(registry, EVENTS_FAILED)).isEqualTo(2.0); - }); - - // The failure released the in-flight slot: the next sealed batch still ships. - uploader.mode = RecordingUploader.Mode.COMPLETE; - handler.publish(record("{\"e\":3}")); - handler.publish(record("{\"e\":4}")); - - await() - .atMost(AWAIT) - .untilAsserted( - () -> { - assertThat(counter(registry, FLUSHED)).isEqualTo(2.0); - assertThat(counter(registry, BATCHES_UPLOADED)).isEqualTo(1.0); - }); - } finally { - handler.close(); - } - } - - @Test - void synchronousUploaderThrowIsContained() { - var uploader = new RecordingUploader(); - var registry = new SimpleMeterRegistry(); - var handler = newHandler(uploader, registry, 2, 1_000_000, 10, 1); - try { - uploader.mode = RecordingUploader.Mode.THROW_SYNC; - handler.publish(record("{\"e\":1}")); - handler.publish(record("{\"e\":2}")); - - await() - .atMost(AWAIT) - .untilAsserted( - () -> { - assertThat(counter(registry, BATCHES_FAILED)).isEqualTo(1.0); - assertThat(counter(registry, EVENTS_FAILED)).isEqualTo(2.0); - }); - - uploader.mode = RecordingUploader.Mode.COMPLETE; - handler.publish(record("{\"e\":3}")); - handler.publish(record("{\"e\":4}")); - - await() - .atMost(AWAIT) - .untilAsserted(() -> assertThat(counter(registry, FLUSHED)).isEqualTo(2.0)); - } finally { - handler.close(); - } - } - - // ============================================================ - // Behavior — close - // ============================================================ - - @Test - void closeDrainsRemainderAndClosesUploader() { - var uploader = new RecordingUploader(); - var registry = new SimpleMeterRegistry(); - var handler = newHandler(uploader, registry, 100, 1_000_000, 10, 2); - - handler.publish(record("{\"e\":1}")); - handler.publish(record("{\"e\":2}")); - handler.publish(record("{\"e\":3}")); - handler.close(); - - // close() is synchronous: by the time it returns the drain has settled and counted. - assertThat(uploader.allLines()) - .containsExactlyInAnyOrder("{\"e\":1}", "{\"e\":2}", "{\"e\":3}"); - assertThat(uploader.closed).isTrue(); - assertThat(counter(registry, FLUSHED)).isEqualTo(3.0); - assertThat(counter(registry, DROPPED, "reason", "shutdown")).isZero(); - } - - @Test - void closeTimeoutCountsAbandoned() { - var uploader = new RecordingUploader(); - var registry = new SimpleMeterRegistry(); - var handler = newHandler(uploader, registry, 1, 1_000_000, 10, 1, Duration.ofMillis(200)); - - uploader.mode = RecordingUploader.Mode.HOLD; - handler.publish(record("{\"e\":1}")); - // Wait until the first batch is in flight (and held) so the queued remainder is exact. - await().atMost(AWAIT).untilAsserted(() -> assertThat(uploader.batches).hasSize(1)); - handler.publish(record("{\"e\":2}")); - handler.publish(record("{\"e\":3}")); - - long startNanos = System.nanoTime(); - handler.close(); - Duration elapsed = Duration.ofNanos(System.nanoTime() - startNanos); - - assertThat(elapsed).isLessThan(Duration.ofSeconds(3)); - assertThat(counter(registry, DROPPED, "reason", "shutdown")).isEqualTo(2.0); - assertThat(uploader.closed).isTrue(); - } - - @Test - void closeIsIdempotentAndPublishAfterCloseIsSafe() { - var uploader = new RecordingUploader(); - var registry = new SimpleMeterRegistry(); - var handler = newHandler(uploader, registry, 1, 1_000_000, 10, 1); - - handler.close(); - handler.close(); // JUL Handler.close() contract: idempotent - - // A racing thread may publish after close; it must never throw (JUL handler contract). - handler.publish(record("{\"late\":1}")); - assertThat(counter(registry, OFFERED)).isEqualTo(1.0); - } - - // ============================================================ - // Async pipeline — gate, chain liveness, back-pressure (deterministic, single driver thread) - // ============================================================ - - @Test - void publishNeverBlocksWhenUploaderStalls() { - var uploader = new RecordingUploader(); - var registry = new SimpleMeterRegistry(); - var handler = newHandler(uploader, registry, 1, 1_000_000, 8, 1); - try { - uploader.mode = RecordingUploader.Mode.HOLD; - handler.publish(record("{\"i\":0}")); - // Wait for the single slot to be claimed and its 1-line batch drained: from here the queue - // is empty, the slot is stuck, and every subsequent count is deterministic. - await().atMost(AWAIT).untilAsserted(() -> assertThat(uploader.batches).hasSize(1)); - - long startNanos = System.nanoTime(); - for (int i = 1; i < 50; i++) { - handler.publish(record("{\"i\":" + i + "}")); - } - Duration elapsed = Duration.ofNanos(System.nanoTime() - startNanos); - - assertThat(elapsed).isLessThan(Duration.ofSeconds(2)); - assertThat(counter(registry, OFFERED)).isEqualTo(50.0); - // 1 in flight + 8 buffered; the other 41 dropped without ever blocking the caller. - assertThat(counter(registry, DROPPED, "reason", "capacity")).isEqualTo(41.0); - } finally { - uploader.releaseAllAndComplete(); - handler.close(); - } - } - - @Test - void concurrencyGateCapsParallelUploads() throws Exception { - var uploader = new RecordingUploader(); - var registry = new SimpleMeterRegistry(); - var handler = newHandler(uploader, registry, 1, 1_000_000, 100, 2); - try { - uploader.mode = RecordingUploader.Mode.HOLD; - for (int i = 0; i < 10; i++) { - handler.publish(record("{\"i\":" + i + "}")); - } - - // Both slots claim work; the rest stays queued behind the gate. - await().atMost(AWAIT).untilAsserted(() -> assertThat(uploader.inFlight.get()).isEqualTo(2)); - - // Each release lets exactly the next batch through, one at a time. - for (int expected = 3; expected <= 10; expected++) { - uploader.releaseOne(); - int size = expected; // fresh effectively-final binding for the lambda - await().atMost(AWAIT).untilAsserted(() -> assertThat(uploader.batches).hasSize(size)); - } - - // maxInFlight = RecordingUploader's high-water mark: uploads peaked at the gate cap of 2. - assertThat(uploader.maxInFlight.get()).isEqualTo(2); - assertThat(uploader.allLines()).hasSize(10).doesNotHaveDuplicates(); - } finally { - uploader.releaseAllAndComplete(); - handler.close(); - } - } - - @Test - void settledUploadChainsNextBatchWithoutNewPublish() throws Exception { - var uploader = new RecordingUploader(); - var registry = new SimpleMeterRegistry(); - var handler = newHandler(uploader, registry, 2, 1_000_000, 100, 1); - try { - uploader.mode = RecordingUploader.Mode.HOLD; - for (int i = 1; i <= 6; i++) { - handler.publish(record("{\"i\":" + i + "}")); - } - - // Single slot: exactly one upload starts, the two other sealed batches wait behind it. - // First await = arrival: the one upload has started. Second = the gate holds: batches stays - // at exactly one for 200ms, proving concurrency=1 keeps the other two sealed batches queued. - await().atMost(AWAIT).untilAsserted(() -> assertThat(uploader.batches).hasSize(1)); - await() - .during(Duration.ofMillis(200)) - .atMost(Duration.ofSeconds(2)) - .until(() -> uploader.batches.size() == 1); - assertThat(uploader.batches.get(0).lines()).containsExactly("{\"i\":1}", "{\"i\":2}"); - - // No further publish happens: each settle must chain the next flush on its own. - uploader.releaseOne(); - await().atMost(AWAIT).untilAsserted(() -> assertThat(uploader.batches).hasSize(2)); - assertThat(uploader.batches.get(1).lines()).containsExactly("{\"i\":3}", "{\"i\":4}"); - - uploader.releaseOne(); - await().atMost(AWAIT).untilAsserted(() -> assertThat(uploader.batches).hasSize(3)); - assertThat(uploader.batches.get(2).lines()).containsExactly("{\"i\":5}", "{\"i\":6}"); - - // Settling the last in-flight upload with an empty buffer chains nothing further. - uploader.releaseOne(); - assertThat(uploader.batches).hasSize(3); - } finally { - uploader.releaseAllAndComplete(); - handler.close(); - } - } - - // ============================================================ - // Concurrent producers — accounting invariants under racing publish (interleaving-agnostic) - // ============================================================ - - @Test - void multiProducerNoLossNoDuplication() throws Exception { - var uploader = new RecordingUploader(); - var registry = new SimpleMeterRegistry(); - var handler = newHandler(uploader, registry, 50, 1_000_000_000L, 10_000, 4); - - int threads = 8; - int perThread = 500; - Set published = ConcurrentHashMap.newKeySet(); - runProducers( - threads, - (threadId) -> { - for (int i = 0; i < perThread; i++) { - String line = "{\"t\":" + threadId + ",\"i\":" + i + "}"; - published.add(line); - handler.publish(record(line)); - } - }); - handler.close(); - - List delivered = uploader.allLines(); - assertThat(delivered).hasSize(threads * perThread); - assertThat(new HashSet<>(delivered)).isEqualTo(published); - assertThat(counter(registry, OFFERED)).isEqualTo(threads * perThread); - assertThat(counter(registry, FLUSHED)).isEqualTo(threads * perThread); - assertThat(counter(registry, DROPPED, "reason", "capacity")).isZero(); - assertThat(counter(registry, DROPPED, "reason", "shutdown")).isZero(); - } - - @Test - void overflowAccountingReconciles() throws Exception { - var uploader = new RecordingUploader(); - var registry = new SimpleMeterRegistry(); - // No seal is ever reached (count seal above capacity, byte seal huge): nothing drains while - // producers run, so every line beyond the 64-slot buffer is a deterministic capacity drop. - var handler = newHandler(uploader, registry, 1000, 1_000_000_000L, 64, 4); - - int threads = 4; - int perThread = 500; - Set published = ConcurrentHashMap.newKeySet(); - runProducers( - threads, - (threadId) -> { - for (int i = 0; i < perThread; i++) { - String line = "{\"t\":" + threadId + ",\"i\":" + i + "}"; - published.add(line); - handler.publish(record(line)); - } - }); - - uploader.releaseAllAndComplete(); - handler.close(); - - double offered = counter(registry, OFFERED); - double flushed = counter(registry, FLUSHED); - double droppedCapacity = counter(registry, DROPPED, "reason", "capacity"); - double droppedShutdown = counter(registry, DROPPED, "reason", "shutdown"); - assertThat(offered).isEqualTo(threads * perThread); - assertThat(flushed).isEqualTo(64.0); - assertThat(droppedCapacity).isEqualTo(threads * perThread - 64.0); - assertThat(flushed + droppedCapacity + droppedShutdown).isEqualTo(offered); - - List delivered = uploader.allLines(); - assertThat(delivered).doesNotHaveDuplicates(); - assertThat(published).containsAll(delivered); - } - - /** - * {@code close()} runs concurrently with in-flight {@code publish()} calls in production — a pod - * shutdown doesn't wait for request threads to go quiet first. This drives both at once: 4 - * threads publish flat out while the main thread calls {@code close()} mid-stream, then keeps the - * producers running a bit longer so some publishes land after close() too. - * - *

Expected: no publish ever throws (the handler must stay safe under this race), close() - * returns instead of hanging, delivered lines are a duplicate-free subset of what was published, - * and the metrics reconcile as {@code flushed + dropped <= offered} rather than {@code ==}. The - * gap is expected, not a bug: a publish can land after close() takes its final buffer snapshot, - * so that line is neither delivered nor counted as dropped — see {@link BatchedLogUploaderMetrics}'s class - * doc for this same at-most-once slippage. The log line below reports the exact gap each run. - */ - @Test - void closeRacingProducersNeverHangsAndReconciles() throws Exception { - var uploader = new RecordingUploader(); - var registry = new SimpleMeterRegistry(); - var handler = newHandler(uploader, registry, 5, 1_000_000_000L, 1000, 4, Duration.ofSeconds(1)); - - int threads = 4; - // Producers normally exit on the stop flag below. The cap bounds the sad path (a hung close - // never reaches stop.set) so no producer spins forever — executor.shutdown() does not interrupt - // running tasks. Overlap with close() is guaranteed by the published.size() gate below. - int perThreadCap = 200_000; - Set published = ConcurrentHashMap.newKeySet(); - List producerErrors = new CopyOnWriteArrayList<>(); - AtomicBoolean stop = new AtomicBoolean(false); - ExecutorService executor = Executors.newFixedThreadPool(threads); - List> futures = new ArrayList<>(); - for (int t = 0; t < threads; t++) { - int threadId = t; - futures.add( - executor.submit( - () -> { - for (int i = 0; i < perThreadCap && !stop.get(); i++) { - String line = "{\"t\":" + threadId + ",\"i\":" + i + "}"; - published.add(line); - try { - handler.publish(record(line)); - } catch (Throwable error) { - producerErrors.add(error); - return; - } - } - })); - } - - // Land close() deterministically amid in-flight publishes: wait until producers have flooded - // the pipeline (2x the 1000-slot buffer → buffer full, overflow dropping, uploads gated), not a - // wall-clock guess. AWAIT only bounds a stuck ramp-up. - await().atMost(AWAIT).until(() -> published.size() >= 2_000); - handler.close(); - stop.set(true); - for (Future future : futures) { - future.get(AWAIT.toSeconds(), TimeUnit.SECONDS); - } - executor.shutdown(); - - // publish must never throw, close must return. Post-close publishes can settle after close()'s - // final buffer snapshot, uncounted, so accounting reconciles with <=, not == (the log shows - // that gap). - // Plain Set ops keep these checks O(n); AssertJ's containsAll would scan, which is O(n^2). - assertThat(producerErrors).isEmpty(); - List delivered = uploader.allLines(); - Set deliveredSet = new HashSet<>(delivered); - assertThat(deliveredSet).as("delivered lines must not repeat").hasSize(delivered.size()); - assertThat(published.containsAll(deliveredSet)) - .as("every delivered line must have been published") - .isTrue(); - double offered = counter(registry, OFFERED); - double flushed = counter(registry, FLUSHED); - double droppedCapacity = counter(registry, DROPPED, "reason", "capacity"); - double droppedShutdown = counter(registry, DROPPED, "reason", "shutdown"); - double accounted = flushed + droppedCapacity + droppedShutdown; - LOG.info( - "closeRacing reconcile: offered={} flushed={} droppedCapacity={} droppedShutdown={}" - + " accounted={} unaccounted={}", - (long) offered, - (long) flushed, - (long) droppedCapacity, - (long) droppedShutdown, - (long) accounted, - (long) (offered - accounted)); - assertThat(accounted).isLessThanOrEqualTo(offered); - } - - // ============================================================ - // Producer harness - // ============================================================ - - private interface Producer { - void run(int threadId) throws Exception; - } - - /** Runs one producer per thread, released simultaneously, and rethrows any producer failure. */ - private static void runProducers(int threads, Producer producer) throws Exception { - ExecutorService executor = Executors.newFixedThreadPool(threads); - try { - CountDownLatch start = new CountDownLatch(1); - List> futures = new ArrayList<>(); - for (int t = 0; t < threads; t++) { - int threadId = t; - futures.add( - executor.submit( - () -> { - start.await(); - producer.run(threadId); - return null; - })); - } - start.countDown(); - for (Future future : futures) { - future.get(30, TimeUnit.SECONDS); - } - } finally { - executor.shutdown(); - } - } + // + // private static final Logger LOG = LoggerFactory.getLogger(BillingS3LogHandlerTest.class); + // + // private static final Duration AWAIT = Duration.ofSeconds(10); + // private static final Duration NEVER = Duration.ofHours(1); + // private static final Instant T0 = Instant.parse("2026-05-20T14:23:11Z"); + // + // /** + // * First-ever Uni creation in a JVM registers the SmallRye context-propagation provider + // through + // * {@code ContextManagerProvider.instance()}, whose ServiceLoader loop both CAS-races + // concurrent + // * callers and throws "ContextManagerProvider already set" when it discovers a second provider + // — + // * possibly after having registered the first. Racing that from concurrent producer threads + // makes + // * publish() throw. Quarkus registers the provider single-threaded at boot, so only this + // * bare-JUnit JVM needs the deterministic warm-up. + // */ + // @BeforeAll + // static void warmUpMutinyInfrastructure() { + // try { + // io.smallrye.context.SmallRyeContextManagerProvider.getManager(); + // } catch (IllegalStateException alreadySetOrDuplicate) { + // // The provider is registered even when the duplicate-discovery branch throws; either way + // // ContextManagerProvider.INSTANCE is now set and concurrent callers can no longer race + // it. + // } + // Uni.createFrom() + // .item(0) + // .runSubscriptionOn(Infrastructure.getDefaultWorkerPool()) + // .await() + // .atMost(AWAIT); + // } + // + // private static final String OFFERED = "billing.s3.events.offered"; + // private static final String FLUSHED = "billing.s3.events.flushed"; + // private static final String EVENTS_FAILED = "billing.s3.events.failed"; + // private static final String BATCHES_UPLOADED = "billing.s3.batches.uploaded"; + // private static final String BATCHES_FAILED = "billing.s3.batches.failed"; + // private static final String DROPPED = "billing.s3.events.dropped"; + // + // // ============================================================ + // // Fake uploader + // // ============================================================ + // + // /** + // * Programmable {@link AsyncBatchedLogUploader}: records every batch and settles the returned + // Uni + // * per {@link Mode}. Never blocks a caller thread — HOLD parks the completion in {@code held} + // for + // * the test to release explicitly. + // */ + // static final class RecordingUploader implements AsyncBatchedLogUploader { + // enum Mode { + // COMPLETE, + // HOLD, + // FAIL, + // THROW_SYNC + // } + // + // volatile Mode mode = Mode.COMPLETE; + // final List batches = new CopyOnWriteArrayList<>(); + // final BlockingQueue> held = new LinkedBlockingQueue<>(); + // final AtomicInteger inFlight = new AtomicInteger(); + // final AtomicInteger maxInFlight = new AtomicInteger(); + // volatile boolean closed; + // + // @Override + // public Uni upload(BatchedLogBuffer.Batch batch) { + // batches.add(batch); + // if (mode == Mode.THROW_SYNC) { + // throw new RuntimeException("simulated synchronous uploader failure"); + // } + // int now = inFlight.incrementAndGet(); + // maxInFlight.accumulateAndGet(now, Math::max); + // CompletableFuture future = new CompletableFuture<>(); + // future.whenComplete((v, t) -> inFlight.decrementAndGet()); + // switch (mode) { + // case COMPLETE -> future.complete(null); + // case FAIL -> future.completeExceptionally(new RuntimeException("simulated upload + // failure")); + // case HOLD -> held.add(future); + // default -> throw new IllegalStateException("unexpected mode " + mode); + // } + // return Uni.createFrom().completionStage(future); + // } + // + // /** Completes one held upload, waiting for it to exist first. */ + // void releaseOne() throws InterruptedException { + // CompletableFuture future = held.poll(AWAIT.toSeconds(), TimeUnit.SECONDS); + // assertThat(future).as("a held upload to release").isNotNull(); + // future.complete(null); + // } + // + // /** Switches to pass-through and completes everything currently held. */ + // void releaseAllAndComplete() { + // mode = Mode.COMPLETE; + // CompletableFuture future; + // while ((future = held.poll()) != null) { + // future.complete(null); + // } + // } + // + // List allLines() { + // return batches.stream().flatMap(b -> b.lines().stream()).toList(); + // } + // + // @Override + // public void close() { + // closed = true; + // } + // } + // + // // ============================================================ + // // Helpers + // // ============================================================ + // + // private static BillingS3LogHandler newHandler( + // RecordingUploader uploader, + // SimpleMeterRegistry registry, + // int maxEvents, + // long maxBytes, + // int queueCapacity, + // int uploadConcurrency) { + // return newHandler( + // uploader, + // registry, + // maxEvents, + // maxBytes, + // queueCapacity, + // uploadConcurrency, + // Duration.ofSeconds(5)); + // } + // + // private static BillingS3LogHandler newHandler( + // RecordingUploader uploader, + // SimpleMeterRegistry registry, + // int maxEvents, + // long maxBytes, + // int queueCapacity, + // int uploadConcurrency, + // Duration shutdownTimeout) { + // return new BillingS3LogHandler( + // uploader, + // registry, + // maxEvents, + // maxBytes, + // NEVER, + // queueCapacity, + // uploadConcurrency, + // shutdownTimeout); + // } + // + // private static LogRecord record(String message) { + // return record(T0, message); + // } + // + // private static LogRecord record(Instant at, String message) { + // LogRecord logRecord = new LogRecord(Level.INFO, message); + // logRecord.setInstant(at); + // return logRecord; + // } + // + // private static double counter(SimpleMeterRegistry registry, String name, String... tags) { + // return registry.counter(name, tags).count(); + // } + // + // // ============================================================ + // // Behavior — publish and flush triggers + // // ============================================================ + // + // @Test + // void publishIgnoresNullRecordAndBlankLines() { + // var uploader = new RecordingUploader(); + // var registry = new SimpleMeterRegistry(); + // var handler = newHandler(uploader, registry, 1, 1_000_000, 10, 1); + // try { + // handler.publish(null); + // handler.publish(record(null)); + // handler.publish(record(" ")); + // + // assertThat(uploader.batches).isEmpty(); + // assertThat(counter(registry, OFFERED)).isZero(); + // } finally { + // handler.close(); + // } + // } + // + // @Test + // void sealsByCountAndShipsExactBatch() { + // var uploader = new RecordingUploader(); + // var registry = new SimpleMeterRegistry(); + // var handler = newHandler(uploader, registry, 3, 1_000_000, 10, 2); + // try { + // // Enqueue order is publish order for a single producer; event-time order is not (the + // second + // // record is older on purpose, so oldestEventAt must be the min, not the head). + // handler.publish(record(T0.plusSeconds(5), "{\"e\":1}")); + // handler.publish(record(T0, "{\"e\":2}")); + // // Asserts the condition holds for the whole window — i.e. that an async flush did NOT + // happen + // await() + // .during(Duration.ofMillis(200)) + // .atMost(Duration.ofSeconds(2)) + // .until(() -> uploader.batches.isEmpty()); + // + // handler.publish(record(T0.plusSeconds(9), "{\"e\":3}")); + // + // await().atMost(AWAIT).untilAsserted(() -> assertThat(uploader.batches).hasSize(1)); + // var batch = uploader.batches.get(0); + // assertThat(batch.lines()).containsExactly("{\"e\":1}", "{\"e\":2}", "{\"e\":3}"); + // assertThat(batch.oldestEventAt()).isEqualTo(T0); + // await() + // .atMost(AWAIT) + // .untilAsserted( + // () -> { + // assertThat(counter(registry, OFFERED)).isEqualTo(3.0); + // assertThat(counter(registry, FLUSHED)).isEqualTo(3.0); + // assertThat(counter(registry, BATCHES_UPLOADED)).isEqualTo(1.0); + // assertThat(counter(registry, DROPPED, "reason", "capacity")).isZero(); + // }); + // } finally { + // handler.close(); + // } + // } + // + // @Test + // void sealsByBufferedBytes() { + // var uploader = new RecordingUploader(); + // var registry = new SimpleMeterRegistry(); + // // Lines count as length + 1: two 4-char lines hit the 10-byte seal together. + // var handler = newHandler(uploader, registry, 100, 10, 10, 2); + // try { + // handler.publish(record("aaaa")); + // handler.publish(record("bbbb")); + // + // await().atMost(AWAIT).untilAsserted(() -> assertThat(uploader.batches).hasSize(1)); + // assertThat(uploader.batches.get(0).lines()).containsExactly("aaaa", "bbbb"); + // } finally { + // handler.close(); + // } + // } + // + // @Test + // void noShipmentBelowSealUntilAgeTick() { + // var uploader = new RecordingUploader(); + // var registry = new SimpleMeterRegistry(); + // var handler = newHandler(uploader, registry, 100, 1_000_000, 10, 2); + // try { + // handler.publish(record("{\"e\":1}")); + // handler.publish(record("{\"e\":2}")); + // await() + // .during(Duration.ofMillis(200)) + // .atMost(Duration.ofSeconds(2)) + // .until(() -> uploader.batches.isEmpty()); + // + // // Deterministic age trigger: call the tick directly instead of waiting for the scheduler. + // handler.onAgeTick(); + // + // await().atMost(AWAIT).untilAsserted(() -> assertThat(uploader.batches).hasSize(1)); + // assertThat(uploader.batches.get(0).lines()).containsExactly("{\"e\":1}", "{\"e\":2}"); + // } finally { + // handler.close(); + // } + // } + // + // @Test + // void ageTickIsScheduledForReal() { + // var uploader = new RecordingUploader(); + // var registry = new SimpleMeterRegistry(); + // // Raw constructor: newHandler() pins maxAge to NEVER; this is the one test that wants a + // live + // // tick. + // var handler = + // new BillingS3LogHandler( + // uploader, + // registry, + // 100, + // 1_000_000, + // Duration.ofMillis(100), + // 10, + // 2, + // Duration.ofSeconds(5)); + // try { + // handler.publish(record("{\"e\":1}")); + // + // // No seal is reached; only the scheduled fixed-rate tick can ship this line. + // await().atMost(AWAIT).untilAsserted(() -> assertThat(uploader.batches).hasSize(1)); + // assertThat(uploader.batches.get(0).lines()).containsExactly("{\"e\":1}"); + // } finally { + // handler.close(); + // } + // } + // + // // ============================================================ + // // Behavior — failure containment + // // ============================================================ + // + // @Test + // void uploadFailureCountsBatchAndPipelineSurvives() { + // var uploader = new RecordingUploader(); + // var registry = new SimpleMeterRegistry(); + // var handler = newHandler(uploader, registry, 2, 1_000_000, 10, 1); + // try { + // uploader.mode = RecordingUploader.Mode.FAIL; + // handler.publish(record("{\"e\":1}")); + // handler.publish(record("{\"e\":2}")); + // + // await() + // .atMost(AWAIT) + // .untilAsserted( + // () -> { + // assertThat(counter(registry, BATCHES_FAILED)).isEqualTo(1.0); + // assertThat(counter(registry, EVENTS_FAILED)).isEqualTo(2.0); + // }); + // + // // The failure released the in-flight slot: the next sealed batch still ships. + // uploader.mode = RecordingUploader.Mode.COMPLETE; + // handler.publish(record("{\"e\":3}")); + // handler.publish(record("{\"e\":4}")); + // + // await() + // .atMost(AWAIT) + // .untilAsserted( + // () -> { + // assertThat(counter(registry, FLUSHED)).isEqualTo(2.0); + // assertThat(counter(registry, BATCHES_UPLOADED)).isEqualTo(1.0); + // }); + // } finally { + // handler.close(); + // } + // } + // + // @Test + // void synchronousUploaderThrowIsContained() { + // var uploader = new RecordingUploader(); + // var registry = new SimpleMeterRegistry(); + // var handler = newHandler(uploader, registry, 2, 1_000_000, 10, 1); + // try { + // uploader.mode = RecordingUploader.Mode.THROW_SYNC; + // handler.publish(record("{\"e\":1}")); + // handler.publish(record("{\"e\":2}")); + // + // await() + // .atMost(AWAIT) + // .untilAsserted( + // () -> { + // assertThat(counter(registry, BATCHES_FAILED)).isEqualTo(1.0); + // assertThat(counter(registry, EVENTS_FAILED)).isEqualTo(2.0); + // }); + // + // uploader.mode = RecordingUploader.Mode.COMPLETE; + // handler.publish(record("{\"e\":3}")); + // handler.publish(record("{\"e\":4}")); + // + // await() + // .atMost(AWAIT) + // .untilAsserted(() -> assertThat(counter(registry, FLUSHED)).isEqualTo(2.0)); + // } finally { + // handler.close(); + // } + // } + // + // // ============================================================ + // // Behavior — close + // // ============================================================ + // + // @Test + // void closeDrainsRemainderAndClosesUploader() { + // var uploader = new RecordingUploader(); + // var registry = new SimpleMeterRegistry(); + // var handler = newHandler(uploader, registry, 100, 1_000_000, 10, 2); + // + // handler.publish(record("{\"e\":1}")); + // handler.publish(record("{\"e\":2}")); + // handler.publish(record("{\"e\":3}")); + // handler.close(); + // + // // close() is synchronous: by the time it returns the drain has settled and counted. + // assertThat(uploader.allLines()) + // .containsExactlyInAnyOrder("{\"e\":1}", "{\"e\":2}", "{\"e\":3}"); + // assertThat(uploader.closed).isTrue(); + // assertThat(counter(registry, FLUSHED)).isEqualTo(3.0); + // assertThat(counter(registry, DROPPED, "reason", "shutdown")).isZero(); + // } + // + // @Test + // void closeTimeoutCountsAbandoned() { + // var uploader = new RecordingUploader(); + // var registry = new SimpleMeterRegistry(); + // var handler = newHandler(uploader, registry, 1, 1_000_000, 10, 1, Duration.ofMillis(200)); + // + // uploader.mode = RecordingUploader.Mode.HOLD; + // handler.publish(record("{\"e\":1}")); + // // Wait until the first batch is in flight (and held) so the queued remainder is exact. + // await().atMost(AWAIT).untilAsserted(() -> assertThat(uploader.batches).hasSize(1)); + // handler.publish(record("{\"e\":2}")); + // handler.publish(record("{\"e\":3}")); + // + // long startNanos = System.nanoTime(); + // handler.close(); + // Duration elapsed = Duration.ofNanos(System.nanoTime() - startNanos); + // + // assertThat(elapsed).isLessThan(Duration.ofSeconds(3)); + // assertThat(counter(registry, DROPPED, "reason", "shutdown")).isEqualTo(2.0); + // assertThat(uploader.closed).isTrue(); + // } + // + // @Test + // void closeIsIdempotentAndPublishAfterCloseIsSafe() { + // var uploader = new RecordingUploader(); + // var registry = new SimpleMeterRegistry(); + // var handler = newHandler(uploader, registry, 1, 1_000_000, 10, 1); + // + // handler.close(); + // handler.close(); // JUL Handler.close() contract: idempotent + // + // // A racing thread may publish after close; it must never throw (JUL handler contract). + // handler.publish(record("{\"late\":1}")); + // assertThat(counter(registry, OFFERED)).isEqualTo(1.0); + // } + // + // // ============================================================ + // // Async pipeline — gate, chain liveness, back-pressure (deterministic, single driver thread) + // // ============================================================ + // + // @Test + // void publishNeverBlocksWhenUploaderStalls() { + // var uploader = new RecordingUploader(); + // var registry = new SimpleMeterRegistry(); + // var handler = newHandler(uploader, registry, 1, 1_000_000, 8, 1); + // try { + // uploader.mode = RecordingUploader.Mode.HOLD; + // handler.publish(record("{\"i\":0}")); + // // Wait for the single slot to be claimed and its 1-line batch drained: from here the + // queue + // // is empty, the slot is stuck, and every subsequent count is deterministic. + // await().atMost(AWAIT).untilAsserted(() -> assertThat(uploader.batches).hasSize(1)); + // + // long startNanos = System.nanoTime(); + // for (int i = 1; i < 50; i++) { + // handler.publish(record("{\"i\":" + i + "}")); + // } + // Duration elapsed = Duration.ofNanos(System.nanoTime() - startNanos); + // + // assertThat(elapsed).isLessThan(Duration.ofSeconds(2)); + // assertThat(counter(registry, OFFERED)).isEqualTo(50.0); + // // 1 in flight + 8 buffered; the other 41 dropped without ever blocking the caller. + // assertThat(counter(registry, DROPPED, "reason", "capacity")).isEqualTo(41.0); + // } finally { + // uploader.releaseAllAndComplete(); + // handler.close(); + // } + // } + // + // @Test + // void concurrencyGateCapsParallelUploads() throws Exception { + // var uploader = new RecordingUploader(); + // var registry = new SimpleMeterRegistry(); + // var handler = newHandler(uploader, registry, 1, 1_000_000, 100, 2); + // try { + // uploader.mode = RecordingUploader.Mode.HOLD; + // for (int i = 0; i < 10; i++) { + // handler.publish(record("{\"i\":" + i + "}")); + // } + // + // // Both slots claim work; the rest stays queued behind the gate. + // await().atMost(AWAIT).untilAsserted(() -> + // assertThat(uploader.inFlight.get()).isEqualTo(2)); + // + // // Each release lets exactly the next batch through, one at a time. + // for (int expected = 3; expected <= 10; expected++) { + // uploader.releaseOne(); + // int size = expected; // fresh effectively-final binding for the lambda + // await().atMost(AWAIT).untilAsserted(() -> assertThat(uploader.batches).hasSize(size)); + // } + // + // // maxInFlight = RecordingUploader's high-water mark: uploads peaked at the gate cap of 2. + // assertThat(uploader.maxInFlight.get()).isEqualTo(2); + // assertThat(uploader.allLines()).hasSize(10).doesNotHaveDuplicates(); + // } finally { + // uploader.releaseAllAndComplete(); + // handler.close(); + // } + // } + // + // @Test + // void settledUploadChainsNextBatchWithoutNewPublish() throws Exception { + // var uploader = new RecordingUploader(); + // var registry = new SimpleMeterRegistry(); + // var handler = newHandler(uploader, registry, 2, 1_000_000, 100, 1); + // try { + // uploader.mode = RecordingUploader.Mode.HOLD; + // for (int i = 1; i <= 6; i++) { + // handler.publish(record("{\"i\":" + i + "}")); + // } + // + // // Single slot: exactly one upload starts, the two other sealed batches wait behind it. + // // First await = arrival: the one upload has started. Second = the gate holds: batches + // stays + // // at exactly one for 200ms, proving concurrency=1 keeps the other two sealed batches + // queued. + // await().atMost(AWAIT).untilAsserted(() -> assertThat(uploader.batches).hasSize(1)); + // await() + // .during(Duration.ofMillis(200)) + // .atMost(Duration.ofSeconds(2)) + // .until(() -> uploader.batches.size() == 1); + // assertThat(uploader.batches.get(0).lines()).containsExactly("{\"i\":1}", "{\"i\":2}"); + // + // // No further publish happens: each settle must chain the next flush on its own. + // uploader.releaseOne(); + // await().atMost(AWAIT).untilAsserted(() -> assertThat(uploader.batches).hasSize(2)); + // assertThat(uploader.batches.get(1).lines()).containsExactly("{\"i\":3}", "{\"i\":4}"); + // + // uploader.releaseOne(); + // await().atMost(AWAIT).untilAsserted(() -> assertThat(uploader.batches).hasSize(3)); + // assertThat(uploader.batches.get(2).lines()).containsExactly("{\"i\":5}", "{\"i\":6}"); + // + // // Settling the last in-flight upload with an empty buffer chains nothing further. + // uploader.releaseOne(); + // assertThat(uploader.batches).hasSize(3); + // } finally { + // uploader.releaseAllAndComplete(); + // handler.close(); + // } + // } + // + // // ============================================================ + // // Concurrent producers — accounting invariants under racing publish (interleaving-agnostic) + // // ============================================================ + // + // @Test + // void multiProducerNoLossNoDuplication() throws Exception { + // var uploader = new RecordingUploader(); + // var registry = new SimpleMeterRegistry(); + // var handler = newHandler(uploader, registry, 50, 1_000_000_000L, 10_000, 4); + // + // int threads = 8; + // int perThread = 500; + // Set published = ConcurrentHashMap.newKeySet(); + // runProducers( + // threads, + // (threadId) -> { + // for (int i = 0; i < perThread; i++) { + // String line = "{\"t\":" + threadId + ",\"i\":" + i + "}"; + // published.add(line); + // handler.publish(record(line)); + // } + // }); + // handler.close(); + // + // List delivered = uploader.allLines(); + // assertThat(delivered).hasSize(threads * perThread); + // assertThat(new HashSet<>(delivered)).isEqualTo(published); + // assertThat(counter(registry, OFFERED)).isEqualTo(threads * perThread); + // assertThat(counter(registry, FLUSHED)).isEqualTo(threads * perThread); + // assertThat(counter(registry, DROPPED, "reason", "capacity")).isZero(); + // assertThat(counter(registry, DROPPED, "reason", "shutdown")).isZero(); + // } + // + // @Test + // void overflowAccountingReconciles() throws Exception { + // var uploader = new RecordingUploader(); + // var registry = new SimpleMeterRegistry(); + // // No seal is ever reached (count seal above capacity, byte seal huge): nothing drains while + // // producers run, so every line beyond the 64-slot buffer is a deterministic capacity drop. + // var handler = newHandler(uploader, registry, 1000, 1_000_000_000L, 64, 4); + // + // int threads = 4; + // int perThread = 500; + // Set published = ConcurrentHashMap.newKeySet(); + // runProducers( + // threads, + // (threadId) -> { + // for (int i = 0; i < perThread; i++) { + // String line = "{\"t\":" + threadId + ",\"i\":" + i + "}"; + // published.add(line); + // handler.publish(record(line)); + // } + // }); + // + // uploader.releaseAllAndComplete(); + // handler.close(); + // + // double offered = counter(registry, OFFERED); + // double flushed = counter(registry, FLUSHED); + // double droppedCapacity = counter(registry, DROPPED, "reason", "capacity"); + // double droppedShutdown = counter(registry, DROPPED, "reason", "shutdown"); + // assertThat(offered).isEqualTo(threads * perThread); + // assertThat(flushed).isEqualTo(64.0); + // assertThat(droppedCapacity).isEqualTo(threads * perThread - 64.0); + // assertThat(flushed + droppedCapacity + droppedShutdown).isEqualTo(offered); + // + // List delivered = uploader.allLines(); + // assertThat(delivered).doesNotHaveDuplicates(); + // assertThat(published).containsAll(delivered); + // } + // + // /** + // * {@code close()} runs concurrently with in-flight {@code publish()} calls in production — a + // pod + // * shutdown doesn't wait for request threads to go quiet first. This drives both at once: 4 + // * threads publish flat out while the main thread calls {@code close()} mid-stream, then keeps + // the + // * producers running a bit longer so some publishes land after close() too. + // * + // *

Expected: no publish ever throws (the handler must stay safe under this race), close() + // * returns instead of hanging, delivered lines are a duplicate-free subset of what was + // published, + // * and the metrics reconcile as {@code flushed + dropped <= offered} rather than {@code ==}. + // The + // * gap is expected, not a bug: a publish can land after close() takes its final buffer + // snapshot, + // * so that line is neither delivered nor counted as dropped — see {@link + // * BatchedLogUploaderMetrics}'s class doc for this same at-most-once slippage. The log line + // below + // * reports the exact gap each run. + // */ + // @Test + // void closeRacingProducersNeverHangsAndReconciles() throws Exception { + // var uploader = new RecordingUploader(); + // var registry = new SimpleMeterRegistry(); + // var handler = newHandler(uploader, registry, 5, 1_000_000_000L, 1000, 4, + // Duration.ofSeconds(1)); + // + // int threads = 4; + // // Producers normally exit on the stop flag below. The cap bounds the sad path (a hung close + // // never reaches stop.set) so no producer spins forever — executor.shutdown() does not + // interrupt + // // running tasks. Overlap with close() is guaranteed by the published.size() gate below. + // int perThreadCap = 200_000; + // Set published = ConcurrentHashMap.newKeySet(); + // List producerErrors = new CopyOnWriteArrayList<>(); + // AtomicBoolean stop = new AtomicBoolean(false); + // ExecutorService executor = Executors.newFixedThreadPool(threads); + // List> futures = new ArrayList<>(); + // for (int t = 0; t < threads; t++) { + // int threadId = t; + // futures.add( + // executor.submit( + // () -> { + // for (int i = 0; i < perThreadCap && !stop.get(); i++) { + // String line = "{\"t\":" + threadId + ",\"i\":" + i + "}"; + // published.add(line); + // try { + // handler.publish(record(line)); + // } catch (Throwable error) { + // producerErrors.add(error); + // return; + // } + // } + // })); + // } + // + // // Land close() deterministically amid in-flight publishes: wait until producers have + // flooded + // // the pipeline (2x the 1000-slot buffer → buffer full, overflow dropping, uploads gated), + // not a + // // wall-clock guess. AWAIT only bounds a stuck ramp-up. + // await().atMost(AWAIT).until(() -> published.size() >= 2_000); + // handler.close(); + // stop.set(true); + // for (Future future : futures) { + // future.get(AWAIT.toSeconds(), TimeUnit.SECONDS); + // } + // executor.shutdown(); + // + // // publish must never throw, close must return. Post-close publishes can settle after + // close()'s + // // final buffer snapshot, uncounted, so accounting reconciles with <=, not == (the log shows + // // that gap). + // // Plain Set ops keep these checks O(n); AssertJ's containsAll would scan, which is O(n^2). + // assertThat(producerErrors).isEmpty(); + // List delivered = uploader.allLines(); + // Set deliveredSet = new HashSet<>(delivered); + // assertThat(deliveredSet).as("delivered lines must not repeat").hasSize(delivered.size()); + // assertThat(published.containsAll(deliveredSet)) + // .as("every delivered line must have been published") + // .isTrue(); + // double offered = counter(registry, OFFERED); + // double flushed = counter(registry, FLUSHED); + // double droppedCapacity = counter(registry, DROPPED, "reason", "capacity"); + // double droppedShutdown = counter(registry, DROPPED, "reason", "shutdown"); + // double accounted = flushed + droppedCapacity + droppedShutdown; + // LOG.info( + // "closeRacing reconcile: offered={} flushed={} droppedCapacity={} droppedShutdown={}" + // + " accounted={} unaccounted={}", + // (long) offered, + // (long) flushed, + // (long) droppedCapacity, + // (long) droppedShutdown, + // (long) accounted, + // (long) (offered - accounted)); + // assertThat(accounted).isLessThanOrEqualTo(offered); + // } + // + // // ============================================================ + // // Producer harness + // // ============================================================ + // + // private interface Producer { + // void run(int threadId) throws Exception; + // } + // + // /** Runs one producer per thread, released simultaneously, and rethrows any producer failure. + // */ + // private static void runProducers(int threads, Producer producer) throws Exception { + // ExecutorService executor = Executors.newFixedThreadPool(threads); + // try { + // CountDownLatch start = new CountDownLatch(1); + // List> futures = new ArrayList<>(); + // for (int t = 0; t < threads; t++) { + // int threadId = t; + // futures.add( + // executor.submit( + // () -> { + // start.await(); + // producer.run(threadId); + // return null; + // })); + // } + // start.countDown(); + // for (Future future : futures) { + // future.get(30, TimeUnit.SECONDS); + // } + // } finally { + // executor.shutdown(); + // } + // } } diff --git a/src/test/java/io/stargate/sgv2/jsonapi/service/billing/DefaultBillingTest.java b/src/test/java/io/stargate/sgv2/jsonapi/service/billing/DefaultBillingTest.java index ac207bafe7..30b570ced6 100644 --- a/src/test/java/io/stargate/sgv2/jsonapi/service/billing/DefaultBillingTest.java +++ b/src/test/java/io/stargate/sgv2/jsonapi/service/billing/DefaultBillingTest.java @@ -9,6 +9,10 @@ import io.stargate.sgv2.jsonapi.TestConstants; import io.stargate.sgv2.jsonapi.config.BillingConfig; +import io.stargate.sgv2.jsonapi.service.provider.ModelInputType; +import io.stargate.sgv2.jsonapi.service.provider.ModelProvider; +import io.stargate.sgv2.jsonapi.service.provider.ModelType; +import io.stargate.sgv2.jsonapi.service.provider.ModelUsage; import java.time.Instant; import java.util.ArrayList; import java.util.Arrays; @@ -20,11 +24,6 @@ import java.util.logging.Handler; import java.util.logging.LogRecord; import java.util.stream.Stream; - -import io.stargate.sgv2.jsonapi.service.provider.ModelInputType; -import io.stargate.sgv2.jsonapi.service.provider.ModelProvider; -import io.stargate.sgv2.jsonapi.service.provider.ModelType; -import io.stargate.sgv2.jsonapi.service.provider.ModelUsage; import org.junit.jupiter.api.Test; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.Arguments; diff --git a/src/test/java/io/stargate/sgv2/jsonapi/service/billing/S3BatchUploaderTest.java b/src/test/java/io/stargate/sgv2/jsonapi/service/billing/S3BatchUploaderTest.java index 9b62fc3f7e..f83ddd596f 100644 --- a/src/test/java/io/stargate/sgv2/jsonapi/service/billing/S3BatchUploaderTest.java +++ b/src/test/java/io/stargate/sgv2/jsonapi/service/billing/S3BatchUploaderTest.java @@ -1,141 +1,122 @@ package io.stargate.sgv2.jsonapi.service.billing; -import static org.assertj.core.api.Assertions.assertThat; -import static org.assertj.core.api.Assertions.assertThatThrownBy; -import static org.mockito.ArgumentMatchers.any; -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.times; -import static org.mockito.Mockito.verify; -import static org.mockito.Mockito.when; - -import java.nio.charset.StandardCharsets; -import java.time.Duration; -import java.time.Instant; -import java.util.List; -import java.util.Optional; -import java.util.UUID; -import java.util.concurrent.CompletableFuture; -import java.util.regex.Pattern; - -import org.junit.jupiter.api.Test; -import org.mockito.ArgumentCaptor; -import software.amazon.awssdk.core.async.AsyncRequestBody; -import software.amazon.awssdk.services.s3.S3AsyncClient; -import software.amazon.awssdk.services.s3.model.PutObjectRequest; -import software.amazon.awssdk.services.s3.model.PutObjectResponse; - /** * The S3 client is mocked; retries and per-call timeouts live in the client configuration, so * exactly one {@code putObject} per upload is expected here. Real I/O is covered by {@code * BillingS3ExportIntegrationTest}. */ class S3BatchUploaderTest { - - private static final Duration AWAIT = Duration.ofSeconds(5); - private static final Pattern KEY_PATTERN = - Pattern.compile("data-api/2026/05/20/14/23/[0-9a-f-]{36}\\.jsonl"); - private static final String LINE_A = "{\"a\":1}"; - private static final String LINE_B = "{\"b\":2}"; - private static final BatchedLogBuffer.Batch BATCH = - new BatchedLogBuffer.Batch(List.of(LINE_A, LINE_B), Instant.parse("2026-05-20T14:23:11.482Z")); - - private static S3BatchedLogUploader uploader(S3AsyncClient client) { - return new S3BatchedLogUploader(client, "my-bucket"); - } - - private static CompletableFuture ok() { - return CompletableFuture.completedFuture(PutObjectResponse.builder().build()); - } - - // ============================================================ - // object layout — key and body - // ============================================================ - - @Test - void objectKeyUsesPathPrefixAndUtcMinutePathFromTimestamp() { - var id = UUID.fromString("8c0e9b8a-1d3a-4f6b-9c0d-1234567890ab"); - var key = S3BatchedLogUploader.objectKey(Instant.parse("2026-05-20T14:23:11.482Z"), id); - assertThat(key) - .isEqualTo("data-api/2026/05/20/14/23/8c0e9b8a-1d3a-4f6b-9c0d-1234567890ab.jsonl"); - } - - @Test - void toNdjsonJoinsLinesVerbatimWithTrailingNewlines() { - assertThat(S3BatchedLogUploader.toNdjson(List.of(LINE_A, LINE_B))) - .isEqualTo((LINE_A + "\n" + LINE_B + "\n").getBytes(StandardCharsets.UTF_8)); - } - - @Test - void putsTheNdjsonBodyAtATimePartitionedKey() { - S3AsyncClient client = mock(S3AsyncClient.class); - when(client.putObject(any(PutObjectRequest.class), any(AsyncRequestBody.class))) - .thenReturn(ok()); - - uploader(client).upload(BATCH).await().atMost(AWAIT); - - var req = ArgumentCaptor.forClass(PutObjectRequest.class); - var body = ArgumentCaptor.forClass(AsyncRequestBody.class); - verify(client).putObject(req.capture(), body.capture()); - assertThat(req.getValue().bucket()).isEqualTo("my-bucket"); - // The key's minute path comes from the batch's oldestEventAt, not the wall clock. - assertThat(req.getValue().key()).matches(KEY_PATTERN.pattern()); - assertThat(req.getValue().contentType()).isEqualTo("application/x-ndjson"); - assertThat(body.getValue().contentLength()) - .hasValue((long) (LINE_A + "\n" + LINE_B + "\n").getBytes(StandardCharsets.UTF_8).length); - } - - @Test - void eachUploadGetsAFreshKey() { - S3AsyncClient client = mock(S3AsyncClient.class); - when(client.putObject(any(PutObjectRequest.class), any(AsyncRequestBody.class))) - .thenReturn(ok()); - - var uploader = uploader(client); - uploader.upload(BATCH).await().atMost(AWAIT); - uploader.upload(BATCH).await().atMost(AWAIT); - - var req = ArgumentCaptor.forClass(PutObjectRequest.class); - verify(client, times(2)).putObject(req.capture(), any(AsyncRequestBody.class)); - assertThat(req.getAllValues().stream().map(PutObjectRequest::key)).doesNotHaveDuplicates(); - } - - // ============================================================ - // failure — surfaces once; retries belong to the SDK client - // ============================================================ - - @Test - void uploadFailurePropagatesAndPutsExactlyOnceAtThisLayer() { - S3AsyncClient client = mock(S3AsyncClient.class); - when(client.putObject(any(PutObjectRequest.class), any(AsyncRequestBody.class))) - .thenReturn(CompletableFuture.failedFuture(new RuntimeException("simulated S3 failure"))); - - var uploader = uploader(client); - assertThatThrownBy(() -> uploader.upload(BATCH).await().atMost(AWAIT)) - .hasMessageContaining("simulated S3 failure"); - - // retries and per-call timeouts live in the client configuration, so exactly one putObject per - // upload is expected here - verify(client, times(1)).putObject(any(PutObjectRequest.class), any(AsyncRequestBody.class)); - } - - // ============================================================ - // lifecycle + config validation - // ============================================================ - - @Test - void closeClosesTheClient() { - S3AsyncClient client = mock(S3AsyncClient.class); - uploader(client).close(); - verify(client).close(); - } - - @Test - void createRejectsMissingRegionOrBucket() { - assertThatThrownBy(() -> S3BatchedLogUploader.create(" ", "bucket", Optional.empty())) - .isInstanceOf(IllegalArgumentException.class) - .hasMessageContaining("bucket-region"); - assertThatThrownBy(() -> S3BatchedLogUploader.create("us-east-1", null, Optional.empty())) - .isInstanceOf(IllegalArgumentException.class) - .hasMessageContaining("billing.s3.bucket"); - } + // + // private static final Duration AWAIT = Duration.ofSeconds(5); + // private static final Pattern KEY_PATTERN = + // Pattern.compile("data-api/2026/05/20/14/23/[0-9a-f-]{36}\\.jsonl"); + // private static final String LINE_A = "{\"a\":1}"; + // private static final String LINE_B = "{\"b\":2}"; + // private static final BatchedLogBuffer.Batch BATCH = + // new BatchedLogBuffer.Batch( + // List.of(LINE_A, LINE_B), Instant.parse("2026-05-20T14:23:11.482Z")); + // + // private static S3BatchedLogUploader uploader(S3AsyncClient client) { + // return new S3BatchedLogUploader(client, "my-bucket"); + // } + // + // private static CompletableFuture ok() { + // return CompletableFuture.completedFuture(PutObjectResponse.builder().build()); + // } + // + // // ============================================================ + // // object layout — key and body + // // ============================================================ + // + // @Test + // void objectKeyUsesPathPrefixAndUtcMinutePathFromTimestamp() { + // var id = UUID.fromString("8c0e9b8a-1d3a-4f6b-9c0d-1234567890ab"); + // var key = S3BatchedLogUploader.objectKey(Instant.parse("2026-05-20T14:23:11.482Z"), id); + // assertThat(key) + // .isEqualTo("data-api/2026/05/20/14/23/8c0e9b8a-1d3a-4f6b-9c0d-1234567890ab.jsonl"); + // } + // + // @Test + // void toNdjsonJoinsLinesVerbatimWithTrailingNewlines() { + // assertThat(S3BatchedLogUploader.toNdjson(List.of(LINE_A, LINE_B))) + // .isEqualTo((LINE_A + "\n" + LINE_B + "\n").getBytes(StandardCharsets.UTF_8)); + // } + // + // @Test + // void putsTheNdjsonBodyAtATimePartitionedKey() { + // S3AsyncClient client = mock(S3AsyncClient.class); + // when(client.putObject(any(PutObjectRequest.class), any(AsyncRequestBody.class))) + // .thenReturn(ok()); + // + // uploader(client).upload(BATCH).await().atMost(AWAIT); + // + // var req = ArgumentCaptor.forClass(PutObjectRequest.class); + // var body = ArgumentCaptor.forClass(AsyncRequestBody.class); + // verify(client).putObject(req.capture(), body.capture()); + // assertThat(req.getValue().bucket()).isEqualTo("my-bucket"); + // // The key's minute path comes from the batch's oldestEventAt, not the wall clock. + // assertThat(req.getValue().key()).matches(KEY_PATTERN.pattern()); + // assertThat(req.getValue().contentType()).isEqualTo("application/x-ndjson"); + // assertThat(body.getValue().contentLength()) + // .hasValue((long) (LINE_A + "\n" + LINE_B + + // "\n").getBytes(StandardCharsets.UTF_8).length); + // } + // + // @Test + // void eachUploadGetsAFreshKey() { + // S3AsyncClient client = mock(S3AsyncClient.class); + // when(client.putObject(any(PutObjectRequest.class), any(AsyncRequestBody.class))) + // .thenReturn(ok()); + // + // var uploader = uploader(client); + // uploader.upload(BATCH).await().atMost(AWAIT); + // uploader.upload(BATCH).await().atMost(AWAIT); + // + // var req = ArgumentCaptor.forClass(PutObjectRequest.class); + // verify(client, times(2)).putObject(req.capture(), any(AsyncRequestBody.class)); + // assertThat(req.getAllValues().stream().map(PutObjectRequest::key)).doesNotHaveDuplicates(); + // } + // + // // ============================================================ + // // failure — surfaces once; retries belong to the SDK client + // // ============================================================ + // + // @Test + // void uploadFailurePropagatesAndPutsExactlyOnceAtThisLayer() { + // S3AsyncClient client = mock(S3AsyncClient.class); + // when(client.putObject(any(PutObjectRequest.class), any(AsyncRequestBody.class))) + // .thenReturn(CompletableFuture.failedFuture(new RuntimeException("simulated S3 + // failure"))); + // + // var uploader = uploader(client); + // assertThatThrownBy(() -> uploader.upload(BATCH).await().atMost(AWAIT)) + // .hasMessageContaining("simulated S3 failure"); + // + // // retries and per-call timeouts live in the client configuration, so exactly one putObject + // per + // // upload is expected here + // verify(client, times(1)).putObject(any(PutObjectRequest.class), + // any(AsyncRequestBody.class)); + // } + // + // // ============================================================ + // // lifecycle + config validation + // // ============================================================ + // + // @Test + // void closeClosesTheClient() { + // S3AsyncClient client = mock(S3AsyncClient.class); + // uploader(client).close(); + // verify(client).close(); + // } + // + // @Test + // void createRejectsMissingRegionOrBucket() { + // assertThatThrownBy(() -> S3BatchedLogUploader.create(" ", "bucket", Optional.empty())) + // .isInstanceOf(IllegalArgumentException.class) + // .hasMessageContaining("bucket-region"); + // assertThatThrownBy(() -> S3BatchedLogUploader.create("us-east-1", null, Optional.empty())) + // .isInstanceOf(IllegalArgumentException.class) + // .hasMessageContaining("billing.s3.bucket"); + // } } diff --git a/src/test/java/io/stargate/sgv2/jsonapi/util/MockClock.java b/src/test/java/io/stargate/sgv2/jsonapi/util/MockClock.java new file mode 100644 index 0000000000..a08e4baa0b --- /dev/null +++ b/src/test/java/io/stargate/sgv2/jsonapi/util/MockClock.java @@ -0,0 +1,62 @@ +package io.stargate.sgv2.jsonapi.util; + +import java.time.*; +import java.util.concurrent.atomic.AtomicReference; + +/** Implementation of the Java Clock that can be used to control time for time dependant tests. */ +public class MockClock extends Clock { + private final Instant startedAt; + private final AtomicReference now; + private final ZoneId zone; + + public MockClock() { + this(Instant.now(), ZoneId.systemDefault()); + } + + public MockClock(MockClock other) { + this(other.startedAt, other.zone); + } + + private MockClock(Instant now, ZoneId zone) { + this.now = new AtomicReference<>(now); + this.startedAt = now; + this.zone = zone; + } + + public Instant startedAt() { + return startedAt; + } + + public MockClock nextSecond() { + return addSeconds(1); + } + + public MockClock addSeconds(int seconds) { + return advance(Duration.ofSeconds(seconds)); + } + + public MockClock advance(Duration amount) { + now.updateAndGet(current -> current.plus(amount)); + return this; + } + + public MockClock setInstant(Instant instant) { + now.set(instant); + return this; + } + + @Override + public ZoneId getZone() { + return zone; + } + + @Override + public Clock withZone(ZoneId zone) { + return new MockClock(now.get(), zone); + } + + @Override + public Instant instant() { + return now.get(); + } +} From 6bef90eee243c4f3f37583adb307c84e3682bc6d Mon Sep 17 00:00:00 2001 From: Aaron Morton Date: Wed, 2 Sep 2026 11:01:42 +1200 Subject: [PATCH 64/65] changes from review --- .../sgv2/jsonapi/metrics/MetricsBase.java | 4 +++- .../service/billing/BatchedLogBuffer.java | 23 +++++++++++-------- .../service/billing/BillingS3LogHandler.java | 10 ++++---- .../service/billing/BatchedLogBufferTest.java | 4 ++-- 4 files changed, 24 insertions(+), 17 deletions(-) diff --git a/src/main/java/io/stargate/sgv2/jsonapi/metrics/MetricsBase.java b/src/main/java/io/stargate/sgv2/jsonapi/metrics/MetricsBase.java index 76461da3e8..c58b2a9cc6 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/metrics/MetricsBase.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/metrics/MetricsBase.java @@ -41,7 +41,9 @@ protected Gauge newGauge(String name, Supplier func) { // no null checks in the builder below Objects.requireNonNull(func, "func must not be null"); - return Gauge.builder(fullName(name), func).register(meterRegistry); + return Gauge.builder(fullName(name), func) + .strongReference(true) // is set in builder() above just being explicit + .register(meterRegistry); } protected Timer newTimer(String name) { diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BatchedLogBuffer.java b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BatchedLogBuffer.java index eed43a3174..105d5b5945 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BatchedLogBuffer.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BatchedLogBuffer.java @@ -75,8 +75,8 @@ public class BatchedLogBuffer { * @param maxBatchBytes Maximum numbers of bytes in a batch, when the buffer has more than this * many entries a new batch is made available which may contain more than this many bytes. The * batch will have many maxBatchBytes if there is a single log record that is bigger. - * @param maxBatchAge Maximum age any log record should have in the buffer before a new batch is - * available. When a batch is triggered from max age the batch is filled, even if the other + * @param maxBatchAge Maximum age the head log record should have in the buffer before a new batch + * is available. When a batch is triggered from max age the batch is filled, even if the other * messages have not reached their max age. * @param capacity Total number of log records to buffer. Beyond this called to {@link * #offer(LogRecord)} will fail to add the message. @@ -108,9 +108,7 @@ public class BatchedLogBuffer { this.maxBatchBytes = maxBatchBytes; this.maxBatchAge = maxBatchAge; this.capacity = capacity; - this.metrics = Objects.requireNonNull(metrics, "billingMetrics must not be null"); - this.metrics.registerBuffer(this); this.clock = clock == null ? DEFAULT_CLOCK : clock; if (this.clock != DEFAULT_CLOCK) { @@ -119,6 +117,9 @@ public class BatchedLogBuffer { } // must be concurrent to handle multiple threads this.queue = new ArrayBlockingQueue<>(capacity); + + // just to be safe, register after queue created incase metrics are scrapped + this.metrics.registerBuffer(this); } /** @@ -216,12 +217,14 @@ public Batch nextBatch(boolean drainFully) { return null; } - LOGGER.info( - "nextBatch() - next batch created, reason:{}, batchLines.size:{}, batchBytes:{}, oldestEventAt: {}", - batchReason, - batchLines.size(), - batchBytes, - oldestEventAt); + if (LOGGER.isDebugEnabled()) { + LOGGER.debug( + "nextBatch() - next batch created, reason:{}, batchLines.size:{}, batchBytes:{}, oldestEventAt: {}", + batchReason, + batchLines.size(), + batchBytes, + oldestEventAt); + } return new Batch(batchReason, batchLines, batchBytes, oldestEventAt, clock); } diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3LogHandler.java b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3LogHandler.java index e1058524bc..76633fe105 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3LogHandler.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3LogHandler.java @@ -12,11 +12,13 @@ import org.slf4j.LoggerFactory; /** - * A Logging handler designed to be used wioth the Billing system. It accpets billing event log - * messges, batches them, and then sends to S3. + * A Logging handler designed to be used with the Billing system. It accepts billing event log + * messages, batches them, and then sends to S3. * - *

See {@link BillingS3HandlerInstaller} for setup. // AI SLOP BELOW JUL handler that turns - * {@code billing.events} log lines into batched S3 objects. + *

See {@link BillingS3HandlerInstaller} for setup. + * + *

// AI SLOP BELOW JUL handler that turns {@code billing.events} log lines into batched S3 + * objects. * *

Division of labor: {@link BatchedLogBuffer} decides when a batch seals, {@link * AsyncBatchedLogUploader} decides what an S3 object looks like, and this class decides when diff --git a/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BatchedLogBufferTest.java b/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BatchedLogBufferTest.java index 7baf16ffa2..94eefc9f5e 100644 --- a/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BatchedLogBufferTest.java +++ b/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BatchedLogBufferTest.java @@ -19,7 +19,7 @@ import java.util.logging.LogRecord; import java.util.stream.IntStream; import java.util.stream.Stream; -import org.junit.Test; +import org.junit.jupiter.api.Test; /** * Unit tests for {@link BatchedLogBuffer}\ @@ -415,7 +415,7 @@ public void nextBatchTriggerDrain() { // 2nd - drainFully - should get a partial batch var batch2 = fixture.assertNextBatch("nextBatchTriggerDrain() - 2nd - partial batch", true); - assertThat(batch1.reason()) + assertThat(batch2.reason()) .as( "nextBatchTriggerDrain() - 2nd - reason is " + BatchedLogBuffer.BillingBatchReason.DRAINING) From de803181c3eb416ca985af618c1ed86f931576c6 Mon Sep 17 00:00:00 2001 From: Aaron Morton Date: Fri, 4 Sep 2026 15:50:24 +1200 Subject: [PATCH 65/65] WIP LogHandler and testing mostly done. still needs decisions on what to do when upload fails. Installer setup to do basic starting of the uploader. --- .../jsonapi/config/BillingS3ExportConfig.java | 26 +- .../billing/AsyncBatchedLogUploader.java | 5 +- .../billing/BillingS3HandlerInstaller.java | 38 +- .../service/billing/BillingS3LogHandler.java | 223 ------ .../billing/BillingUploadingLogHandler.java | 345 ++++++++ .../service/billing/S3BatchedLogUploader.java | 2 +- .../service/billing/BatchedLogBufferTest.java | 409 ++-------- .../billing/BillingS3LogHandlerTest.java | 740 ------------------ .../service/billing/BillingTestBase.java | 484 ++++++++++++ .../BillingUploadingLogHandlerTest.java | 309 ++++++++ 10 files changed, 1228 insertions(+), 1353 deletions(-) delete mode 100644 src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3LogHandler.java create mode 100644 src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingUploadingLogHandler.java delete mode 100644 src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3LogHandlerTest.java create mode 100644 src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingTestBase.java create mode 100644 src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingUploadingLogHandlerTest.java diff --git a/src/main/java/io/stargate/sgv2/jsonapi/config/BillingS3ExportConfig.java b/src/main/java/io/stargate/sgv2/jsonapi/config/BillingS3ExportConfig.java index ef2d7ba3a4..0f10a2b880 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/config/BillingS3ExportConfig.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/config/BillingS3ExportConfig.java @@ -28,21 +28,33 @@ public interface BillingS3ExportConfig { Optional endpointOverride(); /** */ - @WithDefault("2048") - int maxEventsPerBatch(); + @WithDefault("5000") + int maxBatchSize(); /** * Max bytes to include in a batch, NOTE: if a single event is bigger than this it will be sent in * a batch still. 2097152 == 2 MB */ @WithDefault("2097152") - long maxBytesPerBatch(); + long maxBatchBytes(); /** Age flush period: buffered events are shipped at least this often */ - @WithDefault("PT30S") - Duration maxAge(); + @WithDefault("PT60S") + Duration maxBatchAge(); - /** Bound on buffered events; beyond it new lines are dropped. */ - @WithDefault("10000") + /** + * Bound on buffered events; beyond it new lines are dropped. 5,000 events per batch, so set to + * 5,000 * 10 = 50,000 up to 20 MB + */ + @WithDefault("50000") int queueCapacity(); + + @WithDefault("PT60S") + Duration uploadSleepDuration(); + + @WithDefault("PT30S") + Duration uploaderSafetyDeadline(); + + @WithDefault("PT30S") + Duration uploadShutdownDeadline(); } diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/billing/AsyncBatchedLogUploader.java b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/AsyncBatchedLogUploader.java index b08631966a..6a568a6f13 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/service/billing/AsyncBatchedLogUploader.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/AsyncBatchedLogUploader.java @@ -20,9 +20,8 @@ default void close() {} /** * Result of the upload call. * - * @param success true if the operation succeeded, false otherwise. - * @param throwable The throwable associated with an error state. * @param batch The batch that was uploaded, or attempted to be uploaded. + * @param throwable The throwable associated with an error state. */ - record UploadResult(boolean success, Throwable throwable, BatchedLogBuffer.Batch batch) {} + record UploadResult(BatchedLogBuffer.Batch batch, Throwable throwable) {} } diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3HandlerInstaller.java b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3HandlerInstaller.java index b975262e04..50f2ea681c 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3HandlerInstaller.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3HandlerInstaller.java @@ -3,6 +3,7 @@ import io.micrometer.core.instrument.MeterRegistry; import io.quarkus.runtime.ShutdownEvent; import io.quarkus.runtime.StartupEvent; +import io.smallrye.mutiny.infrastructure.Infrastructure; import io.stargate.sgv2.jsonapi.config.BillingS3ExportConfig; import io.stargate.sgv2.jsonapi.metrics.BatchedLogBufferMetrics; import io.stargate.sgv2.jsonapi.metrics.BatchedLogUploaderMetrics; @@ -13,9 +14,9 @@ import org.slf4j.LoggerFactory; /** - * Attaches a {@link BillingS3LogHandler} to the {@code billing.events} logger at startup (when - * {@link BillingS3ExportConfig#enabled()} is {@code true}) and removes + closes it on shutdown for - * a graceful drain. + * Attaches a {@link BillingUploadingLogHandler} to the {@code billing.events} logger at startup + * (when {@link BillingS3ExportConfig#enabled()} is {@code true}) and removes + closes it on + * shutdown for a graceful drain. */ @ApplicationScoped public class BillingS3HandlerInstaller { @@ -29,7 +30,7 @@ public class BillingS3HandlerInstaller { private final BillingS3ExportConfig config; private final MeterRegistry meterRegistry; - private volatile BillingS3LogHandler handler; + private volatile BillingUploadingLogHandler handler; @Inject public BillingS3HandlerInstaller(BillingS3ExportConfig config, MeterRegistry meterRegistry) { @@ -40,10 +41,10 @@ public BillingS3HandlerInstaller(BillingS3ExportConfig config, MeterRegistry met void onStart(@Observes StartupEvent event) { if (!config.enabled()) { - LOGGER.info("Billing S3 export disabled"); + LOGGER.info("onStart() - S3 export disabled"); return; } - LOGGER.info("Billing S3 export enabled"); + LOGGER.info("onStart() - S3 export enabled"); // Fail-loud: invalid billing S3 config throws here, aborting application startup. var uploader = @@ -52,24 +53,31 @@ void onStart(@Observes StartupEvent event) { config.bucket(), config.endpointOverride().orElse(null), new BatchedLogUploaderMetrics(meterRegistry, METRICS_PREFIX)); - LOGGER.info("Billing is using uploader: {}", uploader); + LOGGER.info("onStart() - using uploader: {}", uploader); var buffer = new BatchedLogBuffer( - config.maxEventsPerBatch(), - config.maxBytesPerBatch(), - config.maxAge(), + config.maxBatchSize(), + config.maxBatchBytes(), + config.maxBatchAge(), config.queueCapacity(), new BatchedLogBufferMetrics(meterRegistry, METRICS_PREFIX)); - LOGGER.info("Billing is using log buffer: {}", buffer); - this.handler = new BillingS3LogHandler(buffer, uploader); + LOGGER.info("onStart() - using log buffer: {}", buffer); + + this.handler = + new BillingUploadingLogHandler( + buffer, + uploader, + config.uploadSleepDuration(), + config.uploaderSafetyDeadline(), + config.uploadShutdownDeadline()); + LOGGER.info("onStart() - using uploader: {}", uploader); - // TODO: LOGGER NAME SHOULD BE IN CONFIG Logger.getLogger(BILLING_LOGGER_NAME).addHandler(this.handler); LOGGER.info( - "Billing has attached BillingS3LogHandler to the logger named: {}", BILLING_LOGGER_NAME); + "onStart() - attached log handler to logger. BILLING_LOGGER_NAME: {}", BILLING_LOGGER_NAME); - // TODO: XXXX call start on the thread. + Infrastructure.getDefaultWorkerPool().execute(this.handler::startUploading); } void onStop(@Observes ShutdownEvent event) { diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3LogHandler.java b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3LogHandler.java deleted file mode 100644 index 76633fe105..0000000000 --- a/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3LogHandler.java +++ /dev/null @@ -1,223 +0,0 @@ -package io.stargate.sgv2.jsonapi.service.billing; - -import com.google.common.annotations.VisibleForTesting; -import io.smallrye.mutiny.Uni; -import java.time.Duration; -import java.util.concurrent.CountDownLatch; -import java.util.concurrent.TimeUnit; -import java.util.concurrent.atomic.AtomicBoolean; -import java.util.logging.Handler; -import java.util.logging.LogRecord; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; - -/** - * A Logging handler designed to be used with the Billing system. It accepts billing event log - * messages, batches them, and then sends to S3. - * - *

See {@link BillingS3HandlerInstaller} for setup. - * - *

// AI SLOP BELOW JUL handler that turns {@code billing.events} log lines into batched S3 - * objects. - * - *

Division of labor: {@link BatchedLogBuffer} decides when a batch seals, {@link - * AsyncBatchedLogUploader} decides what an S3 object looks like, and this class decides when - * uploads run — the flush triggers (seal on publish, age tick, drain on close), the - * upload-concurrency gate, and metrics. - * - *

Delivery is at-most-once by design: publish never waits for queue capacity, full buffers drop - * new lines, and close drains best-effort within {@code shutdownTimeout}. - */ -public final class BillingS3LogHandler extends Handler { - - // Logger for this handler, not the destination we are sending events to. - private static final Logger LOGGER = LoggerFactory.getLogger(BillingS3LogHandler.class); - - /** - * Duration the upload thread sleeps between uploading. After all available batches have been - * uploaded, sleeps for this long waiting for the {@link #wakeupSignal} - */ - private static final long UPLOAD_SLEEP_MS = 1000; - - /** - * When signaled this object wakes up the uploading thread to immediately get to work. Used as - * part of the close mechanism to trigger uploading to complete. We do not signal the upload - * everytime a producer calls {@link #publish(LogRecord)}. - */ - private final Object wakeupSignal = new Object(); - - /** - * When true means the Handler has been closed via {@link #close()} and it will silently drop any - * further calls to publish log entries. This also cauese the upload thread to empty the queue - */ - private final AtomicBoolean isClosed = new AtomicBoolean(false); - - /** - * Started at 1 and then decremented in {@link #startUploading()} when it exists so we know we - * have finished uploading. - */ - private final CountDownLatch uploadingFinished = new CountDownLatch(0); - - private final AsyncBatchedLogUploader uploader; - private final BatchedLogBuffer batchedLogBuffer; - - @VisibleForTesting - BillingS3LogHandler(BatchedLogBuffer batchedLogBuffer, AsyncBatchedLogUploader uploader) { - - this.batchedLogBuffer = batchedLogBuffer; - this.uploader = uploader; - } - - private static Duration requirePositive(Duration value, String property) { - if (value == null || value.isNegative() || value.isZero()) { - throw new IllegalArgumentException( - "stargate.jsonapi.billing.s3." + property + " must be > 0 (was " + value + ")"); - } - return value; - } - - // ============================================================ - // Overrides for java.util.logging.Handler - // ============================================================ - - @Override - public void publish(LogRecord record) { - - // Sanity check - if (record == null) { - return; - } - - if (isClosed.get()) { - LOGGER.warn("publish() - called when closed, dropping record:{}", record); - } - - // buffer handles metrics - if (!batchedLogBuffer.offer(record)) { - if (LOGGER.isTraceEnabled()) { - LOGGER.trace("publish() - dropped record:{}", record); - } - } - // flushing runs every second so nothing more to do - } - - @Override - public void flush() { - // wakeup the upload thread to send eveything it can. - // NOTE: this will only drain the buffer if isClosed() is true - notifyUploading(); - } - - /** - * Drains what remains through the normal flush pipeline, bounded by {@code shutdownTimeout}. The - * budget only bites when S3 is already failing: it converts a silent SIGKILL into a logged count - * of abandoned events and lets the rest of shutdown proceed. - */ - @Override - public void close() { - - // mark as closed to stop accepting further events and tell the upload thread - // to drain the bugger/ - isClosed.set(true); - flush(); - - try { - // waiting for the uploading thead to signal it has sent all events in the buffer - if (!uploadingFinished.await(30, TimeUnit.SECONDS)) { - LOGGER.warn("close() - Billing upload loop did not stop within 30s, interrupting"); - } - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - LOGGER.warn("close() - Interrupted waiting for billing upload loop to finish"); - } finally { - uploader.close(); - } - } - - // ============================================================ - // Flush pipeline - // ============================================================ - - /** Called on a worker thread to start uploading log records. */ - void startUploading() { - - BatchedLogBuffer.Batch batch; - try { - while (true) { - - synchronized (wakeupSignal) { - // if the handler is closed we do not want to go to sleep again because it is closing - // down. - if (!isClosed.get()) { - try { - // waiting will release the synchronized monitor - wakeupSignal.wait(1000); - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - return; - } - } - } - - // Get the next batches, if isClosed is true then we want to drain all events - // which may mean creating a batch when we do not have a full one. - - while ((batch = batchedLogBuffer.nextBatch(isClosed.get())) != null) { - - // calling await() on the Uni from deferBatch causes the uni to start running - // indefinitely() is bounded by the 10s ifNoItem() timeout inside deferBatch(), - // and failures are recovered there, so this neither hangs nor throws. - deferBatch(batch).await().indefinitely(); - } - - if (isClosed.get()) { - // Handler is closing down, time to get out of this crazy loop - break; - } - } - } finally { - // record if there are any abandoned events - if (!batchedLogBuffer.isEmpty()) { - LOGGER.warn( - "start() - finished with abandoned billing events, billingQueue.size():{} ", - batchedLogBuffer.size()); - } - - // reset the latch so the call at close() can exit. - uploadingFinished.countDown(); - } - } - - /** - * Signals to the uploading thread that it should wakeup and do some work. - * - *

The uploading thread runs repeated checks for new batches, this is only needed to wakeup as - * part of closing - */ - private void notifyUploading() { - synchronized (wakeupSignal) { - wakeupSignal.notifyAll(); - } - } - - /** - * Creates a Uni that will upload the provided batch. - * - *

As a deferred Uni it does not do any work until something pulls the item, so the caller (see - * startUploading()) starts the work and can decide to wait etc. - * - * @param batch - * @return - */ - private Uni deferBatch(BatchedLogBuffer.Batch batch) { - - // upload() is called at subscription, not when this method returns. - // deferred also converts a synchronous throw from upload() into a Uni failure. - - return Uni.createFrom() - .deferred(() -> uploader.upload(batch)) - .ifNoItem() - .after(Duration.ofSeconds(10)) - .fail(); - } -} diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingUploadingLogHandler.java b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingUploadingLogHandler.java new file mode 100644 index 0000000000..9934dc9468 --- /dev/null +++ b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingUploadingLogHandler.java @@ -0,0 +1,345 @@ +package io.stargate.sgv2.jsonapi.service.billing; + +import com.google.common.annotations.VisibleForTesting; +import io.smallrye.mutiny.Uni; +import java.time.Duration; +import java.util.Objects; +import java.util.concurrent.Semaphore; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.logging.Handler; +import java.util.logging.LogRecord; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * A Logging handler designed to be used with the Billing system. It accepts billing event log + * messages, batches them, and then sends to S3. + * + *

See {@link BillingS3HandlerInstaller} for setup. + * + *

// AI SLOP BELOW JUL handler that turns {@code billing.events} log lines into batched S3 + * objects. + * + *

Division of labor: {@link BatchedLogBuffer} decides when a batch seals, {@link + * AsyncBatchedLogUploader} decides what an S3 object looks like, and this class decides when + * uploads run — the flush triggers (seal on publish, age tick, drain on close), the + * upload-concurrency gate, and metrics. + * + *

Delivery is at-most-once by design: publish never waits for queue capacity, full buffers drop + * new lines, and close drains best-effort within {@code shutdownTimeout}. + */ +public final class BillingUploadingLogHandler extends Handler { + + // Logger for this handler, not the destination we are sending events to. + private static final Logger LOGGER = LoggerFactory.getLogger(BillingUploadingLogHandler.class); + + /** + * When true means the Handler has been closed via {@link #close()} and it will silently drop any + * further calls to publish log entries. This also cause the upload thread to empty the buffer + */ + private final AtomicBoolean isClosed = new AtomicBoolean(false); + + /** + * Disposable permitting system for forcing wakeup in the uploading thread. startUploading() will + * tryAcquire() but because the permit count is 0 will always timeout, this is the timeout to wake + * and check buffer. When we want to force wakeup, e.g. flush(), we call release() that means any + * tryAcquire() returns and decrements count to 0. Resetting back to initial state. Because the + * wakeup permit lasts until tryAcquire it removes race conditions that could happen when + * flush()/notify() on an object lands before the upload thread is in wait() - if we used + * Object.notify() and .wait() + */ + private final Semaphore wakeupPermit = new Semaphore(0); + + /** + * There is only 1 permit for the upload process to be runnning. When {@link #startUploading()} + * starts it takes the permit, gives it back when the function exits (after {@link #close()}. + * close() uses this to make sure uploading has finished. + */ + private final Semaphore uploadPermit = new Semaphore(1); + + private final AsyncBatchedLogUploader uploader; + private final BatchedLogBuffer buffer; + private final Duration uploadSleepDuration; + private final Duration uploaderSafetyDeadline; + private final Duration uploadShutdownDeadline; + + /** See {@link BillingS3HandlerInstaller} */ + BillingUploadingLogHandler( + BatchedLogBuffer buffer, + AsyncBatchedLogUploader uploader, + Duration uploadSleepDuration, + Duration uploaderSafetyDeadline, + Duration uploadShutdownDeadline) { + + this.buffer = Objects.requireNonNull(buffer, "buffer must not be null"); + this.uploader = Objects.requireNonNull(uploader, "uploader must not be null"); + this.uploadSleepDuration = + Objects.requireNonNull(uploadSleepDuration, "uploadSleepDuration must not be null"); + this.uploaderSafetyDeadline = + Objects.requireNonNull(uploaderSafetyDeadline, "uploaderSafetyDeadline must not be null"); + this.uploadShutdownDeadline = + Objects.requireNonNull(uploadShutdownDeadline, "uploadShutdownDeadline must not be null"); + } + + /** + * WARNING - sets the flag for closing but does not run the full close. just here for testing how + * uploading wakes up when flush called. + */ + @VisibleForTesting + void unsafeClose() { + LOGGER.warn("WARNING - unsafeClose() called, must only be used in testing"); + isClosed.set(true); + } + + /** + * WARNING - acquires the upload permit, this stops the startUpload() function and close() from + * working normally. For testing only. + */ + @VisibleForTesting + void unsafeAcquireUploadPermit() { + LOGGER.warn("WARNING - unsafeAcquireUploadPermit() called, must only be used in testing"); + uploadPermit.acquireUninterruptibly(); + } + + // ============================================================ + // Overrides for java.util.logging.Handler + // ============================================================ + + /** + * Buffers and then published the record to S3. + * + * @param record description of the log event. A null record is silently ignored and is not + * published + */ + @Override + public void publish(LogRecord record) { + + // Sanity check + if (record == null) { + return; + } + + if (isClosed.get()) { + LOGGER.warn("publish() - called when closed, dropping record:{}", record); + return; + } + + // buffer handles metrics + if (!buffer.offer(record)) { + if (LOGGER.isDebugEnabled()) { + LOGGER.debug( + "publish() - buffer.offer() rejected, dropping record:{}", record.getMessage()); + } + } else if (LOGGER.isTraceEnabled()) { + LOGGER.trace("publish() - buffer.offer() accepted, record:{}", record.getMessage()); + } + } + + /** + * Wakes up the uploading thread to check the buffer for batches. + * + *

This will only drain the buffer fully (i.e. including partial batches) if {@link #close()} + * is called or {@link #isClosed} is set. + */ + @Override + public void flush() { + maybeTrace("flush() - called"); + notifyUploading(); + } + + /** + * Closes the LogHandler so that it will drop any records sent to {@link #publish(LogRecord)} and + * drain the buffer fully to send all batches to S3. + */ + @Override + public void close() { + + LOGGER.info( + "closing() - marking handler closed, flushing, and waiting for uploads to complete. uploadShutdownDeadline:{}", + uploadShutdownDeadline); + + // mark as closed to stop accepting further events and tell the upload thread + // to drain the buffer fully. + isClosed.set(true); + flush(); + + try { + // Check uploading is not running by trying to get the single uploading permit + // TODO: move timeout to config + if (!uploadPermit.tryAcquire(uploadShutdownDeadline.toMillis(), TimeUnit.MILLISECONDS)) { + LOGGER.warn( + "close() - Failed to get uploading permit, upload failed to stop. uploadShutdownDeadline:{}", + uploadShutdownDeadline); + } else { + uploadPermit.release(); + LOGGER.debug("close() - acquired uploading permit, uploading has completed."); + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + LOGGER.warn("close() - Interrupted waiting for billing upload loop to finish"); + } finally { + uploader.close(); + } + } + + // ============================================================ + // Flush pipeline + // ============================================================ + + /** + * Call this on a worker thread to start uploading, will start a loop of waiting for batches from + * the buffer and uploading them. + */ + void startUploading() { + + LOGGER.info("startUploading() - using buffer:{}, uploader:{}", buffer, uploader); + + boolean hasPermit = false; + BatchedLogBuffer.Batch batch; + try { + if (!(hasPermit = uploadPermit.tryAcquire())) { + throw new IllegalStateException( + "startUploading() - unable to acquire uploadPermit, was function already called?"); + } + + while (true) { + + // if the handler is closed we do not want to go to sleep again because it is closing + // down. + if (!isClosed.get()) { + try { + // waiting will release the synchronized monitor + maybeTrace( + "startUploading() - waiting for wakeupPermit. isClosed:{}, uploadSleepDuration:{}", + isClosed.get(), + uploadSleepDuration); + var acquiredWakePermit = + wakeupPermit.tryAcquire(uploadSleepDuration.toMillis(), TimeUnit.MILLISECONDS); + // is not important if we got a permit to wake, or timed out, just for logging + maybeTrace( + "startUploading() - wakeup permit or timeout, isClosed:{}, acquiredWakePermit:{}", + isClosed.get(), + acquiredWakePermit); + + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + dumpBuffer(); + return; + } + } else { + maybeTrace( + "startUploading() - not waiting for wakeupPermit because isClosed:{}", + isClosed.get()); + } + + // Get the next batches, if isClosed is true then we want to drain all events + // which may mean creating a batch when we do not have a full one. + while ((batch = buffer.nextBatch(isClosed.get())) != null) { + uploadBatch(batch); + } + + if (isClosed.get()) { + // Handler is closing down, time to get out of this crazy loop + break; + } + } + } finally { + // release the uploading permit if we have it, done with the uploading lifestyle + if (hasPermit) { + uploadPermit.release(); + maybeTrace("startUploading() - releasing upload permit"); + } else { + maybeTrace("startUploading() - upload permit was not acquired, not releasing"); + } + } + + if (!buffer.isEmpty()) { + LOGGER.warn( + "startUploading() - finished with abandoned billing events, billingQueue.size():{} ", + buffer.size()); + } + + LOGGER.info( + "startUploading() - stopped uploading using buffer:{}, uploader:{}", buffer, uploader); + } + + /** Adds a permit to the wakeupPermit so the uploading thread will wakeup and do some work. */ + private void notifyUploading() { + wakeupPermit.release(); + } + + private static void maybeTrace(String message, Object... args) { + if (LOGGER.isTraceEnabled()) { + LOGGER.trace(message, args); + } + } + + /** + * Creates a Uni that will upload the provided batch. + * + *

As a deferred Uni it does not do any work until something pulls the item, so the caller (see + * startUploading()) starts the work and can decide to wait etc. + * + * @param batch + * @return + */ + private void uploadBatch(BatchedLogBuffer.Batch batch) { + + LOGGER.info( + "uploadBatch() - starting to upload. uploaderSafetyDeadline:{}, batch:{}", + uploaderSafetyDeadline, + batch); + + // while the uploader should take of all the timeout and retry logic + // as a client of the uploader adding a safety timeout here incase it breaks + + // using deferred so that an error in upload() before it returns the Uni is then + // treated as an error through the Uni pipeline + var uploadResult = + Uni.createFrom() + .deferred(() -> uploader.upload(batch)) + .ifNoItem() + .after(uploaderSafetyDeadline) + .fail() + .onFailure() + .recoverWithItem( + t -> onUploaderFailure(batch, t)) // TimeoutException id deadline exceeded + .await() + .indefinitely(); // the deadline above covers it + + if (uploadResult.throwable() == null) { + onBatchSuccess(uploadResult); + } else { + onBatchFailure(uploadResult); + } + } + + /** + * There was an unhandled error from the uploader(). + * + *

Could be from in upload() before it returned or from running the Uni to do the upload. Just + * map this unhandled back into the UploadResult so we can deal with error in standard way + */ + private AsyncBatchedLogUploader.UploadResult onUploaderFailure( + BatchedLogBuffer.Batch batch, Throwable throwable) { + LOGGER.error( + "onUploaderFailure() - throwable from uploader, adding to UploadResult. batch:{}, throwable:{}", + batch, + throwable.toString()); + return new AsyncBatchedLogUploader.UploadResult(batch, throwable); + } + + private void onBatchSuccess(AsyncBatchedLogUploader.UploadResult uploadResult) { + LOGGER.info("onBatchSuccess() - successfully uploaded batch:{}", uploadResult.batch()); + } + + private void onBatchFailure(AsyncBatchedLogUploader.UploadResult uploadResult) { + LOGGER.error("onBatchFailure() - failed to upload batch:{}", uploadResult.batch()); + } + + /** TODO: dump buffer or a failed batch to regular logs or whatever */ + private void dumpBuffer() {} + + private void dumpBatch() {} +} diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/billing/S3BatchedLogUploader.java b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/S3BatchedLogUploader.java index bcc93fe2ad..2870f5f5d7 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/service/billing/S3BatchedLogUploader.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/S3BatchedLogUploader.java @@ -147,7 +147,7 @@ public Uni upload(BatchedLogBuffer.Batch batch) { resp.eTag(), String.valueOf(resp.sdkHttpResponse().statusCode())); } - return new UploadResult(success, failure, batch); + return new UploadResult(batch, failure); }); } diff --git a/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BatchedLogBufferTest.java b/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BatchedLogBufferTest.java index 94eefc9f5e..4af4868e34 100644 --- a/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BatchedLogBufferTest.java +++ b/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BatchedLogBufferTest.java @@ -7,17 +7,14 @@ import static org.mockito.Mockito.*; import io.stargate.sgv2.jsonapi.metrics.BatchedLogBufferMetrics; -import io.stargate.sgv2.jsonapi.util.MockClock; import java.lang.ref.WeakReference; import java.time.Duration; import java.util.*; import java.util.concurrent.*; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicLong; -import java.util.concurrent.locks.LockSupport; import java.util.logging.Level; import java.util.logging.LogRecord; -import java.util.stream.IntStream; import java.util.stream.Stream; import org.junit.jupiter.api.Test; @@ -27,33 +24,7 @@ *

TODO: out of order log records gets correct oldest metric TODO: TEST a big line bigger than * the max bytes gets through TODO: test metrics using SimpleMeterRegistry */ -public class BatchedLogBufferTest { - - // want the line bytes when lines go into the buffer to be 25 - // template below is 21 bytes - // 3 chars for the index get added in createFixture() - // 1 char added in the buffer calc's for the `\n` to write out - private static final int MESSAGE_LENGTH_IN_BUFFER = 25; - private static final String TEMPLATE_25_CHARS = "Total of 25 chars "; - - private static final int MAX_BATCH_SIZE = 100; - // The number of messages we can fit inside the max bytes setting - private static final int MAX_BATCH_BYTES_NUM_MESSAGES = 20; - private static final int MAX_BATCH_BYTES = - MESSAGE_LENGTH_IN_BUFFER * MAX_BATCH_BYTES_NUM_MESSAGES; - - // How many full batches, tracked by max size, we want to fit in the buffer - private static final int BATCHES_BY_SIZE_PER_CAPACITY = 3; - private static final int BUFFER_CAPACITY = MAX_BATCH_SIZE * BATCHES_BY_SIZE_PER_CAPACITY; - - // number of log records we create for each feature / test - private static final int NUM_RECORDS = BUFFER_CAPACITY * 3; - // when using mock clock, we set the instant for each log record to be 1 "second" - // after the last, so we will create log records with up to - // NUM_RECORDS of seconds past when the clock was started - // used when testing the max age features - private static final Duration MAX_AGE = Duration.ofSeconds(NUM_RECORDS); - private static final Level LOG_LEVEL = Level.INFO; +public class BatchedLogBufferTest extends BillingTestBase { // ********************************************************* // Offer - Producer side of the buffer @@ -63,7 +34,7 @@ public class BatchedLogBufferTest { @Test public void offerFailsAtCapacitySingleThread() { - var fixture = defaultFixture(false); + var fixture = defaultBufferFixture(false); var snapshot = BufferSnapshot.create(fixture); var slice = Slice.to(BUFFER_CAPACITY); @@ -80,15 +51,15 @@ public void offerFailsAtCapacitySingleThread() { @Test public void offerFailsAtCapacityMultiThread() { - var fixture = defaultFixture(false); + var fixture = defaultBufferFixture(false); var snapshot = BufferSnapshot.create(fixture); var slice = Slice.to(BUFFER_CAPACITY); // fill the buffer to capacity from 6 threads calling offer() // auto close will wait for tasks to finish in executor try (var pool = Executors.newFixedThreadPool(6)) { - for (var record : slice.stream(fixture.logRecords).toList()) { - pool.submit(() -> fixture.buffer.offer(record)); + for (var record : slice.stream(fixture.logRecords()).toList()) { + pool.submit(() -> fixture.buffer().offer(record)); } } @@ -105,13 +76,13 @@ public void offerFailsAtCapacityMultiThread() { @Test public void offerDoesNotHoldReferences() { - var fixture = defaultFixture(false); + var fixture = defaultBufferFixture(false); // do not use the records in the fixture, they are held in a list var record = new LogRecord(Level.INFO, "offerDoesNotHoldReferences()"); var ref = new WeakReference<>(record); - fixture.buffer.offer(record); + fixture.buffer().offer(record); record = null; // reference count for the object created for "record" above should now be zero @@ -127,24 +98,24 @@ record = null; @Test public void offerNullRecord() { - var fixture = defaultFixture(false); + var fixture = defaultBufferFixture(false); - assertThatThrownBy(() -> fixture.buffer.offer(null)) + assertThatThrownBy(() -> fixture.buffer().offer(null)) .as("offerNullRecord() null log record is an exception") .isInstanceOf(NullPointerException.class); } @Test public void offerNullOrBlankMessage() { - var fixture = defaultFixture(false); + var fixture = defaultBufferFixture(false); var nullRecord = new LogRecord(Level.INFO, null); var blankRecord = new LogRecord(Level.INFO, " "); - assertThatThrownBy(() -> fixture.buffer.offer(nullRecord)) + assertThatThrownBy(() -> fixture.buffer().offer(nullRecord)) .as("offerNullOrBlankMessage() - null message is an error") .isInstanceOf(IllegalArgumentException.class); - assertThatThrownBy(() -> fixture.buffer.offer(blankRecord)) + assertThatThrownBy(() -> fixture.buffer().offer(blankRecord)) .as("offerNullOrBlankMessage() - blank message is an error") .isInstanceOf(IllegalArgumentException.class); } @@ -158,13 +129,13 @@ public void offerNullOrBlankMessage() { public void nextBatchEmptyBufferNoBatch() { // lock the clock, do not want it to auto advance for batch testing - var fixture = defaultFixture(true); + var fixture = defaultBufferFixture(true); - assertThat(fixture.buffer.nextBatch(false)) + assertThat(fixture.buffer().nextBatch(false)) .as("nextBatchEmptyBufferNoBatch() - drainFully=false, no batch") .isNull(); - assertThat(fixture.buffer.nextBatch(true)) + assertThat(fixture.buffer().nextBatch(true)) .as("nextBatchEmptyBufferNoBatch() - drainFully=true, no batch") .isNull(); } @@ -174,7 +145,7 @@ public void nextBatchEmptyBufferNoBatch() { public void nextBatchBatchProperties() { // lock the clock, do not want it to auto advance for batch testing - var fixture = defaultFixture(true); + var fixture = defaultBufferFixture(true); var slice = Slice.to(BUFFER_CAPACITY); // fill the buffer with all the records it will fit @@ -183,7 +154,7 @@ public void nextBatchBatchProperties() { // keep taking batches and check their properties BatchedLogBuffer.Batch batch; Set batchIds = new HashSet<>(); - while ((batch = fixture.buffer.nextBatch(true)) != null) { + while ((batch = fixture.buffer().nextBatch(true)) != null) { assertThat(batch.id()) .as("nextBatchBatchProperties() - batch ID has not been seen") @@ -197,7 +168,7 @@ public void nextBatchBatchProperties() { .contains("bytes=" + batch.bytes()); } - assertThat(fixture.buffer.isEmpty()) + assertThat(fixture.buffer().isEmpty()) .as("nextBatchBatchProperties() - drained buffer is empty") .isTrue(); } @@ -207,7 +178,7 @@ public void nextBatchBatchProperties() { public void nextBatchMetaUpdatedAfterBatch() { // lock the clock, do not want it to auto advance for batch testing - var fixture = defaultFixture(true); + var fixture = defaultBufferFixture(true); var slice = Slice.to(MAX_BATCH_SIZE); // fill the buffer with 1 batch size and assert metadata @@ -229,7 +200,7 @@ public void nextBatchTriggerMaxSize() { // change so the template is small so does not trigger max bytes // lock the clock, do not want it to auto advance for batch testing var fixture = - createFixture( + Fixture.createFixture( MAX_BATCH_SIZE, MAX_BATCH_BYTES * 100, // big number so never batch because of bytes MAX_AGE, @@ -237,12 +208,15 @@ public void nextBatchTriggerMaxSize() { NUM_RECORDS, LOG_LEVEL, "test-", - true); + true, + false, + false, + false); // Fill to 1 less than max batch size, should be no batch var slice1 = Slice.to(MAX_BATCH_SIZE - 1); fixture.assertOffer("nextBatchMetaUpdatedAfterBatch()", slice1); - var batch1 = fixture.buffer.nextBatch(false); + var batch1 = fixture.buffer().nextBatch(false); assertThat(batch1).as("nextBatchTriggerMaxSize() - < MAX_BATCH_SIZE, no batch").isNull(); // add one more record, should be a batch of MAX_BATCH_SIZE @@ -263,7 +237,7 @@ public void nextBatchTriggerMaxSize() { // add one more record, should be no more batches var slice3 = Slice.slice(MAX_BATCH_SIZE, MAX_BATCH_SIZE + 1); fixture.assertOffer("nextBatchMetaUpdatedAfterBatch() - 3rd", slice3); - var batch3 = fixture.buffer.nextBatch(false); + var batch3 = fixture.buffer().nextBatch(false); assertThat(batch3).as("nextBatchTriggerMaxSize() - 3rd - no batch").isNull(); } @@ -274,12 +248,12 @@ public void nextBatchTriggerMaxBytes() { // default fixture will only fit // the MAX_BATCH_BYTES_NUM_MESSAGES which is less than MAX_SIZE // lock the clock, do not want it to auto advance for batch testing - var fixture = defaultFixture(true); + var fixture = defaultBufferFixture(true); // Fill to 1 message less than max bytes size, should be no batch var slice1 = Slice.to(MAX_BATCH_BYTES_NUM_MESSAGES - 1); fixture.assertOffer("nextBatchTriggerMaxBytes()", slice1); - var batch1 = fixture.buffer.nextBatch(false); + var batch1 = fixture.buffer().nextBatch(false); assertThat(batch1).as("nextBatchTriggerMaxSize() - < MAX_BATCH_BYTES, no batch").isNull(); // add one more , should be a batch of full batch bytes @@ -301,7 +275,7 @@ public void nextBatchTriggerMaxBytes() { // add one more, should be no more batches var slice3 = Slice.slice(MAX_BATCH_BYTES_NUM_MESSAGES, MAX_BATCH_BYTES_NUM_MESSAGES + 1); fixture.assertOffer("nextBatchTriggerMaxBytes() - 3rd", slice3); - var batch3 = fixture.buffer.nextBatch(false); + var batch3 = fixture.buffer().nextBatch(false); assertThat(batch3).as("nextBatchTriggerMaxBytes() - 3rd - no batch").isNull(); } @@ -311,7 +285,7 @@ public void nextBatchTriggerMaxAge() { // lock the clock, do not want it to auto advance for batch testing // NOTE: WE ARE USING THE MOCK CLOCK IN THIS TEST, WE CONTROL TIME - var fixture = defaultFixture(true); + var fixture = defaultBufferFixture(true); // Add only 3 messages, we will not trip size or bytes tigger final int ADDED_RECORDS = 3; @@ -319,7 +293,7 @@ public void nextBatchTriggerMaxAge() { fixture.assertOffer("nextBatchTriggerMaxAge()", slice1); // the clock has not moved, there should be no batch - var batch1 = fixture.buffer.nextBatch(false); + var batch1 = fixture.buffer().nextBatch(false); assertThat(batch1).as("nextBatchTriggerMaxAge() - clock as not moved, no batch").isNull(); // Every LogRecord created in fixture has an instanceAt of 1 second after the previous @@ -333,9 +307,9 @@ public void nextBatchTriggerMaxAge() { // Sanity check, before getting the batch check that only the first log record is MAX_AGE // checking all this junk did what I think int i = 0; - var peekedBuffer = fixture.buffer.peekBuffer(); + var peekedBuffer = fixture.buffer().peekBuffer(); for (var peekEntry : peekedBuffer) { - var entryAge = fixture.buffer.entryAge(peekEntry); + var entryAge = fixture.buffer().entryAge(peekEntry); if (i == 0) { assertThat(entryAge) .as("nextBatchTriggerMaxAge() - clock moved, first entry should be MAX_AGE old") @@ -361,13 +335,13 @@ public void nextBatchTriggerMaxAge() { .isEqualTo(BatchedLogBuffer.BillingBatchReason.MAX_AGE_EXCEEDED); // should have drained all the messages, even if they were not too old - assertThat(fixture.buffer.size()) + assertThat(fixture.buffer().size()) .as("nextBatchTriggerMaxAge() - 2nd batch buffer, size") .isEqualTo(0); - assertThat(fixture.buffer.isEmpty()) + assertThat(fixture.buffer().isEmpty()) .as("nextBatchTriggerMaxAge() - 2nd batch buffer, isEmpty") .isTrue(); - assertThat(fixture.buffer.queuedBytes()) + assertThat(fixture.buffer().queuedBytes()) .as("nextBatchTriggerMaxAge() - 2nd batch buffer, bytes") .isEqualTo(0); @@ -378,12 +352,12 @@ public void nextBatchTriggerMaxAge() { .isEqualTo(ADDED_RECORDS); assertThat(batch2.oldestEventAt()) .as("nextBatchTriggerMaxAge() - 2nd batch buffer, oldest event expected") - .isEqualTo(fixture.logRecords.getFirst().getInstant()); + .isEqualTo(fixture.logRecords().getFirst().getInstant()); // add one more record, should be no more batches var slice3 = Slice.slice(ADDED_RECORDS, ADDED_RECORDS + 1); fixture.assertOffer("nextBatchTriggerMaxAge() - 3rd", slice3); - var batch3 = fixture.buffer.nextBatch(false); + var batch3 = fixture.buffer().nextBatch(false); assertThat(batch3).as("nextBatchTriggerMaxAge() - 3rd - no batch").isNull(); } @@ -395,7 +369,7 @@ public void nextBatchTriggerMaxAge() { public void nextBatchTriggerDrain() { // lock the clock, do not want it to auto advance for batch testing - var fixture = defaultFixture(true); + var fixture = defaultBufferFixture(true); // Fill so we have 1 full batch and 1 partial batch var PARTIAL_BATCH_SIZE = 10; @@ -425,7 +399,7 @@ public void nextBatchTriggerDrain() { .isEqualTo(PARTIAL_BATCH_SIZE); // 3rs - drainFully - no more batch - var batch3 = fixture.buffer.nextBatch(true); + var batch3 = fixture.buffer().nextBatch(true); assertThat(batch3).as("nextBatchTriggerMaxBytes() - 3rd - no batch").isNull(); } @@ -438,7 +412,7 @@ public void nextBatchTriggerDrain() { @Test public void multiThreadedProducerConsumer() { - var fixture = defaultFixture(false); + var fixture = defaultBufferFixture(false); // Setup a Consumer thread, it will keep running until we set consumerShutdown var normalBatches = new ArrayList(); @@ -454,7 +428,7 @@ public void multiThreadedProducerConsumer() { while (!consumerShutdown.get()) { BatchedLogBuffer.Batch consumerNormalBatch; // drainFully=false - because not trying to shutdown - while ((consumerNormalBatch = fixture.buffer.nextBatch(false)) != null) { + while ((consumerNormalBatch = fixture.buffer().nextBatch(false)) != null) { normalBatches.add(consumerNormalBatch); // fake that we do some work with the batch, e.g. upload it threadSleep(50); @@ -465,7 +439,7 @@ public void multiThreadedProducerConsumer() { // now into the shutdown mode, so drainFully=true to empty the buffer BatchedLogBuffer.Batch consumerShutdownBatch; - while ((consumerShutdownBatch = fixture.buffer.nextBatch(true)) != null) { + while ((consumerShutdownBatch = fixture.buffer().nextBatch(true)) != null) { shutdownBatches.add(consumerShutdownBatch); // fake that we do some work with the batch, e.g. upload it threadSleep(50); @@ -484,7 +458,7 @@ public void multiThreadedProducerConsumer() { var producerHalfwayLatch = new CountDownLatch(NUM_PRODUCER_THREADS); try (var pool = Executors.newFixedThreadPool(NUM_PRODUCER_THREADS, threadFactory)) { - for (var record : slice.stream(fixture.logRecords).toList()) { + for (var record : slice.stream(fixture.logRecords()).toList()) { // Append the thread name to the log record for debugging // this will break the config at top of class about how many messages per batch @@ -492,7 +466,7 @@ public void multiThreadedProducerConsumer() { () -> { record.setMessage( record.getMessage() + " - THREAD " + Thread.currentThread().getName()); - fixture.buffer.offer(record); + fixture.buffer().offer(record); if ((producedCount.incrementAndGet() >= (slice.size() / 2)) && (!consumerShutdown.get())) { @@ -617,297 +591,4 @@ public void testConstructor() { .contains("maxBatchAge=PT1S") .contains("size=0"); } - - // ********************************************************* - // Scaffold - // ********************************************************* - - private static void threadSleep(long millis) { - LockSupport.parkNanos(Duration.ofMillis(millis).toNanos()); - } - - private static void waitOnLatch(CountDownLatch latch) { - try { - latch.await(); - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - throw new RuntimeException("interrupted waiting on latch", e); - } - } - - /** Default fixture with config from the top of class */ - private Fixture defaultFixture(boolean mockBufferClock) { - return createFixture( - MAX_BATCH_SIZE, - MAX_BATCH_BYTES, - MAX_AGE, - BUFFER_CAPACITY, - NUM_RECORDS, - LOG_LEVEL, - TEMPLATE_25_CHARS, - mockBufferClock); - } - - /** Create fixture, creates LogRecords that can be used to add to the buffer */ - private Fixture createFixture( - int maxBatchSize, - long maxBytes, - Duration maxAge, - int queueCapacity, - int numLogRecords, - Level logLevel, - String logRecordTemplate, - boolean mockBufferClock) { - - // Make sure to initialize the mock clock before creating the log messages - // so they are always after the start of the clock. - var mockClock = mockBufferClock ? new MockClock() : null; - - // fork the clock, we are going to use clockForRecords when creating the records - // and will advance it 1 second for each record, the original mockClock is for - // the buffer to use, so we let the test advance that - var clockForRecords = mockClock == null ? null : new MockClock(mockClock); - - var logRecords = - IntStream.range(0, numLogRecords) - .mapToObj(i -> logRecordTemplate + String.format("%03d", i)) - .map( - s -> { - var record = new LogRecord(logLevel, s); - if (clockForRecords != null) { - record.setInstant(clockForRecords.instant()); - clockForRecords.nextSecond(); - } - return record; - }) - .toList(); - - var metrics = mock(BatchedLogBufferMetrics.class); - - var buffer = - new BatchedLogBuffer( - maxBatchSize, - maxBytes, - maxAge, - queueCapacity, - metrics, - mockBufferClock ? mockClock : BatchedLogBuffer.DEFAULT_CLOCK); - - return new Fixture( - maxBatchSize, maxBytes, maxAge, queueCapacity, logRecords, buffer, metrics, mockClock); - } - - /** A slice of a list, `from` is inclusive, `to` is exclusive */ - record Slice(int from, int to) { - - public Stream stream(List list) { - return list.stream().skip(from).limit(to - from); - } - - public int size() { - return to - from; - } - - public static Slice to(int to) { - return new Slice(0, to); - } - - public static Slice from(int from) { - return new Slice(from, Integer.MAX_VALUE); - } - - public static Slice slice(int from, int to) { - return new Slice(from, to); - } - } - - /** - * Snapshot of the metadata (size etc) for the buffer, that can be used to compare how the buffer - * metadata has changed - */ - record BufferSnapshot( - boolean isEmpty, int size, long queuedBytes, int remainingCapacity, Fixture fixture) { - - static BufferSnapshot create(Fixture fixture) { - // reset the counters for calls to metrics - clearInvocations(fixture.metrics); - return new BufferSnapshot( - fixture.buffer.isEmpty(), - fixture.buffer.size(), - fixture.buffer.queuedBytes(), - fixture.buffer.remainingCapacity(), - fixture); - } - - /** - * Assert that the current metadata values for the buffer are the values in the snapshot PLUS - * the log records that were added by the Slice. - */ - void assertAll(String desc, Slice slice, boolean inOrder) { - assertBufferMetadata(desc, slice); - assertBufferItems(desc, slice, inOrder); - } - - /** - * Assert that the current metadata values for the buffer are the values in the snapshot MINUS - * the buffer entries that were removed in the batch - */ - void assertAll(String desc, BatchedLogBuffer.Batch batch) { - assertBufferMetadata(desc, batch); - assertBufferItems(desc, batch); - } - - /** current buffer metadata = snapshot + slice */ - void assertBufferMetadata(String desc, Slice slice) { - - if (slice.size() == 0) { - assertThat(fixture.buffer.isEmpty()) - .as(desc + " - isEmpty no change after empty slice") - .isEqualTo(isEmpty()); - } else { - assertThat(fixture.buffer.isEmpty()) - .as(desc + " - isEmpty false after non empty slice") - .isEqualTo(false); - } - - assertThat(fixture.buffer.size()) - .as(desc + " - post buffer size increased by slice") - .isEqualTo(size() + slice.size()); - - verify( - fixture.metrics, - times(slice.size()).description(desc + "metrics called for every offer")) - .offered(); - - long addedBytes = 0; - for (var record : slice.stream(fixture.logRecords).toList()) { - addedBytes += BatchedLogBuffer.Entry.lineBytes(record.getMessage()); - } - - assertThat(fixture.buffer.queuedBytes()) - .as(desc + " - post buffer bytes increased by slice") - .isEqualTo(queuedBytes + addedBytes); - } - - /** current buffer metadata = snapshot - batch */ - void assertBufferMetadata(String desc, BatchedLogBuffer.Batch batch) { - - assertThat(fixture.buffer.size()) - .as(desc + " - buffer size decreased by batch size") - .isEqualTo(size() - batch.size()); - - assertThat(fixture.buffer.queuedBytes()) - .as(desc + " - buffer bytes size decreased by batch bytes") - .isEqualTo(queuedBytes - batch.bytes()); - } - - /** - * current buffer items contain items from slice inOrder - if we expect items in buffer to match - * order of the fixture - */ - void assertBufferItems(String desc, Slice slice, boolean inOrder) { - - var bufferItems = fixture.buffer.peekBuffer(); - - int i = slice.from() > bufferItems.size() ? 0 : slice.from(); - for (var record : slice.stream(fixture.logRecords).toList()) { - - if (inOrder) { - assertThat(record.getMessage()) - .as(desc + " - buffer items at position match exactly pos: " + i) - .isEqualTo(bufferItems.get(i++).line()); - } else { - - var entry = new BatchedLogBuffer.Entry(record.getInstant(), record.getMessage()); - assertThat(bufferItems) - .as(desc + " - buffer items contains entry: " + entry) - .contains(entry); - } - } - } - - /** current buffer items contain NONE of items in batch */ - void assertBufferItems(String desc, BatchedLogBuffer.Batch batch) { - - var peekedBuffer = fixture.buffer.peekBuffer(); - - for (var batchString : batch.lines()) { - - var found = peekedBuffer.stream().anyMatch(entry -> entry.line().equals(batchString)); - assertThat(found) - .as(desc + " - line from batch no longer in buffer: " + batchString) - .isFalse(); - } - } - } - - /** - * Tracks the config of the buffer, the buffer, the data we can use for each test to add to - * buffer, etc. - * - *

See {@link #defaultFixture(boolean)} - */ - record Fixture( - int maxBatchSize, - long maxBytes, - Duration maxAge, - int queueCapacity, - List logRecords, - BatchedLogBuffer buffer, - BatchedLogBufferMetrics metrics, - MockClock clock) { - - /** Assert the buffer is full, and so offer() fails */ - void assertBufferFull(String desc, int index) { - - // although the next log record is wafer-thin, it is too much for Mr Creosote - assertThat(buffer().offer(logRecords.get(index))).as(desc + " - fail at capacity").isFalse(); - - // Running again to confirm it is still full - assertThat(buffer().offer(logRecords.get(index))) - .as(desc + " - second - fail at capacity") - .isFalse(); - } - - /** - * Offer the log records selected by slice to the buffer, all should work, assert the buffer has - * the items the slice selected - */ - void assertOffer(String desc, Slice slice) { - - var snapshot = BufferSnapshot.create(this); - - for (var record : slice.stream(logRecords).toList()) { - assertThat(buffer.offer(record)).as(desc + " - assertOffer() - offering").isTrue(); - } - - snapshot.assertAll(desc, slice, true); - } - - /** - * Get a batch from the buffer, assert we got a batch that is legal, and assert the buffer has - * changed by the amount of the batch - */ - BatchedLogBuffer.Batch assertNextBatch(String desc, boolean drainFully) { - - var snapshot = BufferSnapshot.create(this); - var batch = buffer.nextBatch(drainFully); - - // assert the batch is what we expected. - assertThat(batch).as(desc + " - assertNextBatch() - batch is not null").isNotNull(); - - assertThat(batch.size()) - .as(desc + " - assertNextBatch() - batch size <= MAX_BATCH_SIZE") - .isLessThanOrEqualTo(maxBatchSize); - // note: it is legal to have a batch bigger than the maxBytes, specialised tests for that - // shoudl only happen when there is a single log record bigger than maxBytes - assertThat(batch.bytes()) - .as(desc + " - assertNextBatch() - batch bytes <= MAX_BATCH_BYTES") - .isLessThanOrEqualTo(maxBytes); - - // assert the buffer updated bookkeeping as we expect - snapshot.assertAll(desc, batch); - return batch; - } - } } diff --git a/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3LogHandlerTest.java b/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3LogHandlerTest.java deleted file mode 100644 index 165a3bd21a..0000000000 --- a/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3LogHandlerTest.java +++ /dev/null @@ -1,740 +0,0 @@ -package io.stargate.sgv2.jsonapi.service.billing; - -/** - * Unit tests for {@link BillingS3LogHandler}: the flush triggers (seal on publish, age tick, drain - * on close), the upload-concurrency gate, failure containment, and the at-most-once accounting - * invariants under concurrent publish. The uploader is a programmable in-memory fake; real S3 I/O - * is covered by {@code BillingS3ExportIntegrationTest}. - */ -class BillingS3LogHandlerTest { - // - // private static final Logger LOG = LoggerFactory.getLogger(BillingS3LogHandlerTest.class); - // - // private static final Duration AWAIT = Duration.ofSeconds(10); - // private static final Duration NEVER = Duration.ofHours(1); - // private static final Instant T0 = Instant.parse("2026-05-20T14:23:11Z"); - // - // /** - // * First-ever Uni creation in a JVM registers the SmallRye context-propagation provider - // through - // * {@code ContextManagerProvider.instance()}, whose ServiceLoader loop both CAS-races - // concurrent - // * callers and throws "ContextManagerProvider already set" when it discovers a second provider - // — - // * possibly after having registered the first. Racing that from concurrent producer threads - // makes - // * publish() throw. Quarkus registers the provider single-threaded at boot, so only this - // * bare-JUnit JVM needs the deterministic warm-up. - // */ - // @BeforeAll - // static void warmUpMutinyInfrastructure() { - // try { - // io.smallrye.context.SmallRyeContextManagerProvider.getManager(); - // } catch (IllegalStateException alreadySetOrDuplicate) { - // // The provider is registered even when the duplicate-discovery branch throws; either way - // // ContextManagerProvider.INSTANCE is now set and concurrent callers can no longer race - // it. - // } - // Uni.createFrom() - // .item(0) - // .runSubscriptionOn(Infrastructure.getDefaultWorkerPool()) - // .await() - // .atMost(AWAIT); - // } - // - // private static final String OFFERED = "billing.s3.events.offered"; - // private static final String FLUSHED = "billing.s3.events.flushed"; - // private static final String EVENTS_FAILED = "billing.s3.events.failed"; - // private static final String BATCHES_UPLOADED = "billing.s3.batches.uploaded"; - // private static final String BATCHES_FAILED = "billing.s3.batches.failed"; - // private static final String DROPPED = "billing.s3.events.dropped"; - // - // // ============================================================ - // // Fake uploader - // // ============================================================ - // - // /** - // * Programmable {@link AsyncBatchedLogUploader}: records every batch and settles the returned - // Uni - // * per {@link Mode}. Never blocks a caller thread — HOLD parks the completion in {@code held} - // for - // * the test to release explicitly. - // */ - // static final class RecordingUploader implements AsyncBatchedLogUploader { - // enum Mode { - // COMPLETE, - // HOLD, - // FAIL, - // THROW_SYNC - // } - // - // volatile Mode mode = Mode.COMPLETE; - // final List batches = new CopyOnWriteArrayList<>(); - // final BlockingQueue> held = new LinkedBlockingQueue<>(); - // final AtomicInteger inFlight = new AtomicInteger(); - // final AtomicInteger maxInFlight = new AtomicInteger(); - // volatile boolean closed; - // - // @Override - // public Uni upload(BatchedLogBuffer.Batch batch) { - // batches.add(batch); - // if (mode == Mode.THROW_SYNC) { - // throw new RuntimeException("simulated synchronous uploader failure"); - // } - // int now = inFlight.incrementAndGet(); - // maxInFlight.accumulateAndGet(now, Math::max); - // CompletableFuture future = new CompletableFuture<>(); - // future.whenComplete((v, t) -> inFlight.decrementAndGet()); - // switch (mode) { - // case COMPLETE -> future.complete(null); - // case FAIL -> future.completeExceptionally(new RuntimeException("simulated upload - // failure")); - // case HOLD -> held.add(future); - // default -> throw new IllegalStateException("unexpected mode " + mode); - // } - // return Uni.createFrom().completionStage(future); - // } - // - // /** Completes one held upload, waiting for it to exist first. */ - // void releaseOne() throws InterruptedException { - // CompletableFuture future = held.poll(AWAIT.toSeconds(), TimeUnit.SECONDS); - // assertThat(future).as("a held upload to release").isNotNull(); - // future.complete(null); - // } - // - // /** Switches to pass-through and completes everything currently held. */ - // void releaseAllAndComplete() { - // mode = Mode.COMPLETE; - // CompletableFuture future; - // while ((future = held.poll()) != null) { - // future.complete(null); - // } - // } - // - // List allLines() { - // return batches.stream().flatMap(b -> b.lines().stream()).toList(); - // } - // - // @Override - // public void close() { - // closed = true; - // } - // } - // - // // ============================================================ - // // Helpers - // // ============================================================ - // - // private static BillingS3LogHandler newHandler( - // RecordingUploader uploader, - // SimpleMeterRegistry registry, - // int maxEvents, - // long maxBytes, - // int queueCapacity, - // int uploadConcurrency) { - // return newHandler( - // uploader, - // registry, - // maxEvents, - // maxBytes, - // queueCapacity, - // uploadConcurrency, - // Duration.ofSeconds(5)); - // } - // - // private static BillingS3LogHandler newHandler( - // RecordingUploader uploader, - // SimpleMeterRegistry registry, - // int maxEvents, - // long maxBytes, - // int queueCapacity, - // int uploadConcurrency, - // Duration shutdownTimeout) { - // return new BillingS3LogHandler( - // uploader, - // registry, - // maxEvents, - // maxBytes, - // NEVER, - // queueCapacity, - // uploadConcurrency, - // shutdownTimeout); - // } - // - // private static LogRecord record(String message) { - // return record(T0, message); - // } - // - // private static LogRecord record(Instant at, String message) { - // LogRecord logRecord = new LogRecord(Level.INFO, message); - // logRecord.setInstant(at); - // return logRecord; - // } - // - // private static double counter(SimpleMeterRegistry registry, String name, String... tags) { - // return registry.counter(name, tags).count(); - // } - // - // // ============================================================ - // // Behavior — publish and flush triggers - // // ============================================================ - // - // @Test - // void publishIgnoresNullRecordAndBlankLines() { - // var uploader = new RecordingUploader(); - // var registry = new SimpleMeterRegistry(); - // var handler = newHandler(uploader, registry, 1, 1_000_000, 10, 1); - // try { - // handler.publish(null); - // handler.publish(record(null)); - // handler.publish(record(" ")); - // - // assertThat(uploader.batches).isEmpty(); - // assertThat(counter(registry, OFFERED)).isZero(); - // } finally { - // handler.close(); - // } - // } - // - // @Test - // void sealsByCountAndShipsExactBatch() { - // var uploader = new RecordingUploader(); - // var registry = new SimpleMeterRegistry(); - // var handler = newHandler(uploader, registry, 3, 1_000_000, 10, 2); - // try { - // // Enqueue order is publish order for a single producer; event-time order is not (the - // second - // // record is older on purpose, so oldestEventAt must be the min, not the head). - // handler.publish(record(T0.plusSeconds(5), "{\"e\":1}")); - // handler.publish(record(T0, "{\"e\":2}")); - // // Asserts the condition holds for the whole window — i.e. that an async flush did NOT - // happen - // await() - // .during(Duration.ofMillis(200)) - // .atMost(Duration.ofSeconds(2)) - // .until(() -> uploader.batches.isEmpty()); - // - // handler.publish(record(T0.plusSeconds(9), "{\"e\":3}")); - // - // await().atMost(AWAIT).untilAsserted(() -> assertThat(uploader.batches).hasSize(1)); - // var batch = uploader.batches.get(0); - // assertThat(batch.lines()).containsExactly("{\"e\":1}", "{\"e\":2}", "{\"e\":3}"); - // assertThat(batch.oldestEventAt()).isEqualTo(T0); - // await() - // .atMost(AWAIT) - // .untilAsserted( - // () -> { - // assertThat(counter(registry, OFFERED)).isEqualTo(3.0); - // assertThat(counter(registry, FLUSHED)).isEqualTo(3.0); - // assertThat(counter(registry, BATCHES_UPLOADED)).isEqualTo(1.0); - // assertThat(counter(registry, DROPPED, "reason", "capacity")).isZero(); - // }); - // } finally { - // handler.close(); - // } - // } - // - // @Test - // void sealsByBufferedBytes() { - // var uploader = new RecordingUploader(); - // var registry = new SimpleMeterRegistry(); - // // Lines count as length + 1: two 4-char lines hit the 10-byte seal together. - // var handler = newHandler(uploader, registry, 100, 10, 10, 2); - // try { - // handler.publish(record("aaaa")); - // handler.publish(record("bbbb")); - // - // await().atMost(AWAIT).untilAsserted(() -> assertThat(uploader.batches).hasSize(1)); - // assertThat(uploader.batches.get(0).lines()).containsExactly("aaaa", "bbbb"); - // } finally { - // handler.close(); - // } - // } - // - // @Test - // void noShipmentBelowSealUntilAgeTick() { - // var uploader = new RecordingUploader(); - // var registry = new SimpleMeterRegistry(); - // var handler = newHandler(uploader, registry, 100, 1_000_000, 10, 2); - // try { - // handler.publish(record("{\"e\":1}")); - // handler.publish(record("{\"e\":2}")); - // await() - // .during(Duration.ofMillis(200)) - // .atMost(Duration.ofSeconds(2)) - // .until(() -> uploader.batches.isEmpty()); - // - // // Deterministic age trigger: call the tick directly instead of waiting for the scheduler. - // handler.onAgeTick(); - // - // await().atMost(AWAIT).untilAsserted(() -> assertThat(uploader.batches).hasSize(1)); - // assertThat(uploader.batches.get(0).lines()).containsExactly("{\"e\":1}", "{\"e\":2}"); - // } finally { - // handler.close(); - // } - // } - // - // @Test - // void ageTickIsScheduledForReal() { - // var uploader = new RecordingUploader(); - // var registry = new SimpleMeterRegistry(); - // // Raw constructor: newHandler() pins maxAge to NEVER; this is the one test that wants a - // live - // // tick. - // var handler = - // new BillingS3LogHandler( - // uploader, - // registry, - // 100, - // 1_000_000, - // Duration.ofMillis(100), - // 10, - // 2, - // Duration.ofSeconds(5)); - // try { - // handler.publish(record("{\"e\":1}")); - // - // // No seal is reached; only the scheduled fixed-rate tick can ship this line. - // await().atMost(AWAIT).untilAsserted(() -> assertThat(uploader.batches).hasSize(1)); - // assertThat(uploader.batches.get(0).lines()).containsExactly("{\"e\":1}"); - // } finally { - // handler.close(); - // } - // } - // - // // ============================================================ - // // Behavior — failure containment - // // ============================================================ - // - // @Test - // void uploadFailureCountsBatchAndPipelineSurvives() { - // var uploader = new RecordingUploader(); - // var registry = new SimpleMeterRegistry(); - // var handler = newHandler(uploader, registry, 2, 1_000_000, 10, 1); - // try { - // uploader.mode = RecordingUploader.Mode.FAIL; - // handler.publish(record("{\"e\":1}")); - // handler.publish(record("{\"e\":2}")); - // - // await() - // .atMost(AWAIT) - // .untilAsserted( - // () -> { - // assertThat(counter(registry, BATCHES_FAILED)).isEqualTo(1.0); - // assertThat(counter(registry, EVENTS_FAILED)).isEqualTo(2.0); - // }); - // - // // The failure released the in-flight slot: the next sealed batch still ships. - // uploader.mode = RecordingUploader.Mode.COMPLETE; - // handler.publish(record("{\"e\":3}")); - // handler.publish(record("{\"e\":4}")); - // - // await() - // .atMost(AWAIT) - // .untilAsserted( - // () -> { - // assertThat(counter(registry, FLUSHED)).isEqualTo(2.0); - // assertThat(counter(registry, BATCHES_UPLOADED)).isEqualTo(1.0); - // }); - // } finally { - // handler.close(); - // } - // } - // - // @Test - // void synchronousUploaderThrowIsContained() { - // var uploader = new RecordingUploader(); - // var registry = new SimpleMeterRegistry(); - // var handler = newHandler(uploader, registry, 2, 1_000_000, 10, 1); - // try { - // uploader.mode = RecordingUploader.Mode.THROW_SYNC; - // handler.publish(record("{\"e\":1}")); - // handler.publish(record("{\"e\":2}")); - // - // await() - // .atMost(AWAIT) - // .untilAsserted( - // () -> { - // assertThat(counter(registry, BATCHES_FAILED)).isEqualTo(1.0); - // assertThat(counter(registry, EVENTS_FAILED)).isEqualTo(2.0); - // }); - // - // uploader.mode = RecordingUploader.Mode.COMPLETE; - // handler.publish(record("{\"e\":3}")); - // handler.publish(record("{\"e\":4}")); - // - // await() - // .atMost(AWAIT) - // .untilAsserted(() -> assertThat(counter(registry, FLUSHED)).isEqualTo(2.0)); - // } finally { - // handler.close(); - // } - // } - // - // // ============================================================ - // // Behavior — close - // // ============================================================ - // - // @Test - // void closeDrainsRemainderAndClosesUploader() { - // var uploader = new RecordingUploader(); - // var registry = new SimpleMeterRegistry(); - // var handler = newHandler(uploader, registry, 100, 1_000_000, 10, 2); - // - // handler.publish(record("{\"e\":1}")); - // handler.publish(record("{\"e\":2}")); - // handler.publish(record("{\"e\":3}")); - // handler.close(); - // - // // close() is synchronous: by the time it returns the drain has settled and counted. - // assertThat(uploader.allLines()) - // .containsExactlyInAnyOrder("{\"e\":1}", "{\"e\":2}", "{\"e\":3}"); - // assertThat(uploader.closed).isTrue(); - // assertThat(counter(registry, FLUSHED)).isEqualTo(3.0); - // assertThat(counter(registry, DROPPED, "reason", "shutdown")).isZero(); - // } - // - // @Test - // void closeTimeoutCountsAbandoned() { - // var uploader = new RecordingUploader(); - // var registry = new SimpleMeterRegistry(); - // var handler = newHandler(uploader, registry, 1, 1_000_000, 10, 1, Duration.ofMillis(200)); - // - // uploader.mode = RecordingUploader.Mode.HOLD; - // handler.publish(record("{\"e\":1}")); - // // Wait until the first batch is in flight (and held) so the queued remainder is exact. - // await().atMost(AWAIT).untilAsserted(() -> assertThat(uploader.batches).hasSize(1)); - // handler.publish(record("{\"e\":2}")); - // handler.publish(record("{\"e\":3}")); - // - // long startNanos = System.nanoTime(); - // handler.close(); - // Duration elapsed = Duration.ofNanos(System.nanoTime() - startNanos); - // - // assertThat(elapsed).isLessThan(Duration.ofSeconds(3)); - // assertThat(counter(registry, DROPPED, "reason", "shutdown")).isEqualTo(2.0); - // assertThat(uploader.closed).isTrue(); - // } - // - // @Test - // void closeIsIdempotentAndPublishAfterCloseIsSafe() { - // var uploader = new RecordingUploader(); - // var registry = new SimpleMeterRegistry(); - // var handler = newHandler(uploader, registry, 1, 1_000_000, 10, 1); - // - // handler.close(); - // handler.close(); // JUL Handler.close() contract: idempotent - // - // // A racing thread may publish after close; it must never throw (JUL handler contract). - // handler.publish(record("{\"late\":1}")); - // assertThat(counter(registry, OFFERED)).isEqualTo(1.0); - // } - // - // // ============================================================ - // // Async pipeline — gate, chain liveness, back-pressure (deterministic, single driver thread) - // // ============================================================ - // - // @Test - // void publishNeverBlocksWhenUploaderStalls() { - // var uploader = new RecordingUploader(); - // var registry = new SimpleMeterRegistry(); - // var handler = newHandler(uploader, registry, 1, 1_000_000, 8, 1); - // try { - // uploader.mode = RecordingUploader.Mode.HOLD; - // handler.publish(record("{\"i\":0}")); - // // Wait for the single slot to be claimed and its 1-line batch drained: from here the - // queue - // // is empty, the slot is stuck, and every subsequent count is deterministic. - // await().atMost(AWAIT).untilAsserted(() -> assertThat(uploader.batches).hasSize(1)); - // - // long startNanos = System.nanoTime(); - // for (int i = 1; i < 50; i++) { - // handler.publish(record("{\"i\":" + i + "}")); - // } - // Duration elapsed = Duration.ofNanos(System.nanoTime() - startNanos); - // - // assertThat(elapsed).isLessThan(Duration.ofSeconds(2)); - // assertThat(counter(registry, OFFERED)).isEqualTo(50.0); - // // 1 in flight + 8 buffered; the other 41 dropped without ever blocking the caller. - // assertThat(counter(registry, DROPPED, "reason", "capacity")).isEqualTo(41.0); - // } finally { - // uploader.releaseAllAndComplete(); - // handler.close(); - // } - // } - // - // @Test - // void concurrencyGateCapsParallelUploads() throws Exception { - // var uploader = new RecordingUploader(); - // var registry = new SimpleMeterRegistry(); - // var handler = newHandler(uploader, registry, 1, 1_000_000, 100, 2); - // try { - // uploader.mode = RecordingUploader.Mode.HOLD; - // for (int i = 0; i < 10; i++) { - // handler.publish(record("{\"i\":" + i + "}")); - // } - // - // // Both slots claim work; the rest stays queued behind the gate. - // await().atMost(AWAIT).untilAsserted(() -> - // assertThat(uploader.inFlight.get()).isEqualTo(2)); - // - // // Each release lets exactly the next batch through, one at a time. - // for (int expected = 3; expected <= 10; expected++) { - // uploader.releaseOne(); - // int size = expected; // fresh effectively-final binding for the lambda - // await().atMost(AWAIT).untilAsserted(() -> assertThat(uploader.batches).hasSize(size)); - // } - // - // // maxInFlight = RecordingUploader's high-water mark: uploads peaked at the gate cap of 2. - // assertThat(uploader.maxInFlight.get()).isEqualTo(2); - // assertThat(uploader.allLines()).hasSize(10).doesNotHaveDuplicates(); - // } finally { - // uploader.releaseAllAndComplete(); - // handler.close(); - // } - // } - // - // @Test - // void settledUploadChainsNextBatchWithoutNewPublish() throws Exception { - // var uploader = new RecordingUploader(); - // var registry = new SimpleMeterRegistry(); - // var handler = newHandler(uploader, registry, 2, 1_000_000, 100, 1); - // try { - // uploader.mode = RecordingUploader.Mode.HOLD; - // for (int i = 1; i <= 6; i++) { - // handler.publish(record("{\"i\":" + i + "}")); - // } - // - // // Single slot: exactly one upload starts, the two other sealed batches wait behind it. - // // First await = arrival: the one upload has started. Second = the gate holds: batches - // stays - // // at exactly one for 200ms, proving concurrency=1 keeps the other two sealed batches - // queued. - // await().atMost(AWAIT).untilAsserted(() -> assertThat(uploader.batches).hasSize(1)); - // await() - // .during(Duration.ofMillis(200)) - // .atMost(Duration.ofSeconds(2)) - // .until(() -> uploader.batches.size() == 1); - // assertThat(uploader.batches.get(0).lines()).containsExactly("{\"i\":1}", "{\"i\":2}"); - // - // // No further publish happens: each settle must chain the next flush on its own. - // uploader.releaseOne(); - // await().atMost(AWAIT).untilAsserted(() -> assertThat(uploader.batches).hasSize(2)); - // assertThat(uploader.batches.get(1).lines()).containsExactly("{\"i\":3}", "{\"i\":4}"); - // - // uploader.releaseOne(); - // await().atMost(AWAIT).untilAsserted(() -> assertThat(uploader.batches).hasSize(3)); - // assertThat(uploader.batches.get(2).lines()).containsExactly("{\"i\":5}", "{\"i\":6}"); - // - // // Settling the last in-flight upload with an empty buffer chains nothing further. - // uploader.releaseOne(); - // assertThat(uploader.batches).hasSize(3); - // } finally { - // uploader.releaseAllAndComplete(); - // handler.close(); - // } - // } - // - // // ============================================================ - // // Concurrent producers — accounting invariants under racing publish (interleaving-agnostic) - // // ============================================================ - // - // @Test - // void multiProducerNoLossNoDuplication() throws Exception { - // var uploader = new RecordingUploader(); - // var registry = new SimpleMeterRegistry(); - // var handler = newHandler(uploader, registry, 50, 1_000_000_000L, 10_000, 4); - // - // int threads = 8; - // int perThread = 500; - // Set published = ConcurrentHashMap.newKeySet(); - // runProducers( - // threads, - // (threadId) -> { - // for (int i = 0; i < perThread; i++) { - // String line = "{\"t\":" + threadId + ",\"i\":" + i + "}"; - // published.add(line); - // handler.publish(record(line)); - // } - // }); - // handler.close(); - // - // List delivered = uploader.allLines(); - // assertThat(delivered).hasSize(threads * perThread); - // assertThat(new HashSet<>(delivered)).isEqualTo(published); - // assertThat(counter(registry, OFFERED)).isEqualTo(threads * perThread); - // assertThat(counter(registry, FLUSHED)).isEqualTo(threads * perThread); - // assertThat(counter(registry, DROPPED, "reason", "capacity")).isZero(); - // assertThat(counter(registry, DROPPED, "reason", "shutdown")).isZero(); - // } - // - // @Test - // void overflowAccountingReconciles() throws Exception { - // var uploader = new RecordingUploader(); - // var registry = new SimpleMeterRegistry(); - // // No seal is ever reached (count seal above capacity, byte seal huge): nothing drains while - // // producers run, so every line beyond the 64-slot buffer is a deterministic capacity drop. - // var handler = newHandler(uploader, registry, 1000, 1_000_000_000L, 64, 4); - // - // int threads = 4; - // int perThread = 500; - // Set published = ConcurrentHashMap.newKeySet(); - // runProducers( - // threads, - // (threadId) -> { - // for (int i = 0; i < perThread; i++) { - // String line = "{\"t\":" + threadId + ",\"i\":" + i + "}"; - // published.add(line); - // handler.publish(record(line)); - // } - // }); - // - // uploader.releaseAllAndComplete(); - // handler.close(); - // - // double offered = counter(registry, OFFERED); - // double flushed = counter(registry, FLUSHED); - // double droppedCapacity = counter(registry, DROPPED, "reason", "capacity"); - // double droppedShutdown = counter(registry, DROPPED, "reason", "shutdown"); - // assertThat(offered).isEqualTo(threads * perThread); - // assertThat(flushed).isEqualTo(64.0); - // assertThat(droppedCapacity).isEqualTo(threads * perThread - 64.0); - // assertThat(flushed + droppedCapacity + droppedShutdown).isEqualTo(offered); - // - // List delivered = uploader.allLines(); - // assertThat(delivered).doesNotHaveDuplicates(); - // assertThat(published).containsAll(delivered); - // } - // - // /** - // * {@code close()} runs concurrently with in-flight {@code publish()} calls in production — a - // pod - // * shutdown doesn't wait for request threads to go quiet first. This drives both at once: 4 - // * threads publish flat out while the main thread calls {@code close()} mid-stream, then keeps - // the - // * producers running a bit longer so some publishes land after close() too. - // * - // *

Expected: no publish ever throws (the handler must stay safe under this race), close() - // * returns instead of hanging, delivered lines are a duplicate-free subset of what was - // published, - // * and the metrics reconcile as {@code flushed + dropped <= offered} rather than {@code ==}. - // The - // * gap is expected, not a bug: a publish can land after close() takes its final buffer - // snapshot, - // * so that line is neither delivered nor counted as dropped — see {@link - // * BatchedLogUploaderMetrics}'s class doc for this same at-most-once slippage. The log line - // below - // * reports the exact gap each run. - // */ - // @Test - // void closeRacingProducersNeverHangsAndReconciles() throws Exception { - // var uploader = new RecordingUploader(); - // var registry = new SimpleMeterRegistry(); - // var handler = newHandler(uploader, registry, 5, 1_000_000_000L, 1000, 4, - // Duration.ofSeconds(1)); - // - // int threads = 4; - // // Producers normally exit on the stop flag below. The cap bounds the sad path (a hung close - // // never reaches stop.set) so no producer spins forever — executor.shutdown() does not - // interrupt - // // running tasks. Overlap with close() is guaranteed by the published.size() gate below. - // int perThreadCap = 200_000; - // Set published = ConcurrentHashMap.newKeySet(); - // List producerErrors = new CopyOnWriteArrayList<>(); - // AtomicBoolean stop = new AtomicBoolean(false); - // ExecutorService executor = Executors.newFixedThreadPool(threads); - // List> futures = new ArrayList<>(); - // for (int t = 0; t < threads; t++) { - // int threadId = t; - // futures.add( - // executor.submit( - // () -> { - // for (int i = 0; i < perThreadCap && !stop.get(); i++) { - // String line = "{\"t\":" + threadId + ",\"i\":" + i + "}"; - // published.add(line); - // try { - // handler.publish(record(line)); - // } catch (Throwable error) { - // producerErrors.add(error); - // return; - // } - // } - // })); - // } - // - // // Land close() deterministically amid in-flight publishes: wait until producers have - // flooded - // // the pipeline (2x the 1000-slot buffer → buffer full, overflow dropping, uploads gated), - // not a - // // wall-clock guess. AWAIT only bounds a stuck ramp-up. - // await().atMost(AWAIT).until(() -> published.size() >= 2_000); - // handler.close(); - // stop.set(true); - // for (Future future : futures) { - // future.get(AWAIT.toSeconds(), TimeUnit.SECONDS); - // } - // executor.shutdown(); - // - // // publish must never throw, close must return. Post-close publishes can settle after - // close()'s - // // final buffer snapshot, uncounted, so accounting reconciles with <=, not == (the log shows - // // that gap). - // // Plain Set ops keep these checks O(n); AssertJ's containsAll would scan, which is O(n^2). - // assertThat(producerErrors).isEmpty(); - // List delivered = uploader.allLines(); - // Set deliveredSet = new HashSet<>(delivered); - // assertThat(deliveredSet).as("delivered lines must not repeat").hasSize(delivered.size()); - // assertThat(published.containsAll(deliveredSet)) - // .as("every delivered line must have been published") - // .isTrue(); - // double offered = counter(registry, OFFERED); - // double flushed = counter(registry, FLUSHED); - // double droppedCapacity = counter(registry, DROPPED, "reason", "capacity"); - // double droppedShutdown = counter(registry, DROPPED, "reason", "shutdown"); - // double accounted = flushed + droppedCapacity + droppedShutdown; - // LOG.info( - // "closeRacing reconcile: offered={} flushed={} droppedCapacity={} droppedShutdown={}" - // + " accounted={} unaccounted={}", - // (long) offered, - // (long) flushed, - // (long) droppedCapacity, - // (long) droppedShutdown, - // (long) accounted, - // (long) (offered - accounted)); - // assertThat(accounted).isLessThanOrEqualTo(offered); - // } - // - // // ============================================================ - // // Producer harness - // // ============================================================ - // - // private interface Producer { - // void run(int threadId) throws Exception; - // } - // - // /** Runs one producer per thread, released simultaneously, and rethrows any producer failure. - // */ - // private static void runProducers(int threads, Producer producer) throws Exception { - // ExecutorService executor = Executors.newFixedThreadPool(threads); - // try { - // CountDownLatch start = new CountDownLatch(1); - // List> futures = new ArrayList<>(); - // for (int t = 0; t < threads; t++) { - // int threadId = t; - // futures.add( - // executor.submit( - // () -> { - // start.await(); - // producer.run(threadId); - // return null; - // })); - // } - // start.countDown(); - // for (Future future : futures) { - // future.get(30, TimeUnit.SECONDS); - // } - // } finally { - // executor.shutdown(); - // } - // } -} diff --git a/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingTestBase.java b/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingTestBase.java new file mode 100644 index 0000000000..a506a83db8 --- /dev/null +++ b/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingTestBase.java @@ -0,0 +1,484 @@ +package io.stargate.sgv2.jsonapi.service.billing; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.mockito.Mockito.*; + +import io.stargate.sgv2.jsonapi.metrics.BatchedLogBufferMetrics; +import io.stargate.sgv2.jsonapi.util.MockClock; +import java.time.Duration; +import java.util.List; +import java.util.concurrent.*; +import java.util.concurrent.locks.LockSupport; +import java.util.logging.Level; +import java.util.logging.LogRecord; +import java.util.stream.IntStream; +import java.util.stream.Stream; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** Common code for tests around billing events being uploaded */ +public abstract class BillingTestBase { + + private static final Logger LOGGER = LoggerFactory.getLogger(BillingTestBase.class); + + // ====================================================================== + // BUFFER config and values for the tests to use + // ====================================================================== + + // want the line bytes when lines go into the buffer to be 25 + // template below is 21 bytes + // 3 chars for the index get added in createFixture() + // 1 char added in the buffer calc's for the `\n` to write out + protected static final int MESSAGE_LENGTH_IN_BUFFER = 25; + protected static final String TEMPLATE_25_CHARS = "Total of 25 chars "; + + protected static final int MAX_BATCH_SIZE = 100; + // The number of messages we can fit inside the max bytes setting + protected static final int MAX_BATCH_BYTES_NUM_MESSAGES = 20; + protected static final int MAX_BATCH_BYTES = + MESSAGE_LENGTH_IN_BUFFER * MAX_BATCH_BYTES_NUM_MESSAGES; + + // How many full batches, tracked by max size, we want to fit in the buffer + protected static final int BATCHES_BY_SIZE_PER_CAPACITY = 3; + protected static final int BUFFER_CAPACITY = MAX_BATCH_SIZE * BATCHES_BY_SIZE_PER_CAPACITY; + + // number of log records we create for each feature / test + protected static final int NUM_RECORDS = BUFFER_CAPACITY * 3; + // when using mock clock, we set the instant for each log record to be 1 "second" + // after the last, so we will create log records with up to + // NUM_RECORDS of seconds past when the clock was started + // used when testing the max age features + protected static final Duration MAX_AGE = Duration.ofSeconds(NUM_RECORDS); + protected static final Level LOG_LEVEL = Level.INFO; + + // ====================================================================== + // LOG HANDLER config and values for the tests to use + // ====================================================================== + + // sleep between checking the buffer, long we will normally use flush() to wake + // for tests + protected static final Duration UPLOAD_SLEEP_DURATION = Duration.ofSeconds(60); + protected static final Duration UPLOAD_SLEEP_DURATION_SHORT = Duration.ofMillis(100); + + // upload must complete in this time + protected static final Duration UPLOADER_SAFETY_DEADLINE = Duration.ofSeconds(30); + // close() will wait this long for startUploading() to finish + protected static final Duration UPLOAD_SHUTDOWN_DEADLINE = Duration.ofSeconds(30); + protected static final Duration UPLOAD_SHUTDOWN_DEADLINE_SHORT = Duration.ofSeconds(1); + + protected static void threadSleep(long millis) { + LockSupport.parkNanos(Duration.ofMillis(millis).toNanos()); + } + + protected static void waitOnLatch(CountDownLatch latch) { + try { + latch.await(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException("interrupted waiting on latch", e); + } + } + + // ********************************************************* + // Fixture - used to group the test config and test data we need + // ********************************************************* + + /** Default fixture for buffer tests with config from the top of class */ + Fixture defaultBufferFixture() { + return defaultBufferFixture(false); + } + + Fixture defaultBufferFixture(boolean mockBufferClock) { + return Fixture.createFixture( + MAX_BATCH_SIZE, + MAX_BATCH_BYTES, + MAX_AGE, + BUFFER_CAPACITY, + NUM_RECORDS, + LOG_LEVEL, + TEMPLATE_25_CHARS, + mockBufferClock, + false, + false, + false); + } + + Fixture defaultLogHandlerFixture() { + return defaultLogHandlerFixture(true, false, false); + } + + Fixture defaultLogHandlerFixture( + boolean mockBuffer, boolean shortUploadSleep, boolean shortShutdownDuration) { + return Fixture.createFixture( + MAX_BATCH_SIZE, + MAX_BATCH_BYTES, + MAX_AGE, + BUFFER_CAPACITY, + NUM_RECORDS, + LOG_LEVEL, + TEMPLATE_25_CHARS, + false, + mockBuffer, + shortUploadSleep, + shortShutdownDuration); + } + + /** + * Tracks the config of the buffer, the buffer, the data we can use for each test to add to + * buffer, etc. + * + *

See {@link #defaultBufferFixture(boolean)} + */ + record Fixture( + int maxBatchSize, + long maxBytes, + Duration maxAge, + int queueCapacity, + List logRecords, + BatchedLogBuffer buffer, + BatchedLogBufferMetrics metrics, + BillingUploadingLogHandler logHandler, + AsyncBatchedLogUploader uploader, + MockClock clock) { + + /** Create fixture, creates LogRecords that can be used to add to the buffer */ + static Fixture createFixture( + int maxBatchSize, + long maxBytes, + Duration maxAge, + int queueCapacity, + int numLogRecords, + Level logLevel, + String logRecordTemplate, + boolean mockBufferClock, + boolean mockBuffer, + boolean shortUploadSleepDuration, + boolean shortUploadShutdownDeadline) { + + if (mockBuffer && mockBufferClock) { + throw new IllegalArgumentException("cannot mock the buffer and the buffer clock"); + } + // Make sure to initialize the mock clock before creating the log messages + // so they are always after the start of the clock. + var mockClock = mockBufferClock ? new MockClock() : null; + + // fork the clock, we are going to use clockForRecords when creating the records + // and will advance it 1 second for each record, the original mockClock is for + // the buffer to use, so we let the test advance that + var clockForRecords = mockClock == null ? null : new MockClock(mockClock); + + var logRecords = + IntStream.range(0, numLogRecords) + .mapToObj(i -> logRecordTemplate + String.format("%03d", i)) + .map( + s -> { + var record = new LogRecord(logLevel, s); + if (clockForRecords != null) { + record.setInstant(clockForRecords.instant()); + clockForRecords.nextSecond(); + } + return record; + }) + .toList(); + + var metrics = mock(BatchedLogBufferMetrics.class); + + BatchedLogBuffer buffer; + if (mockBuffer) { + buffer = mock(BatchedLogBuffer.class); + // setup for an empty buffer when calling + when(buffer.offer(any())).thenReturn(true); + when(buffer.isEmpty()).thenReturn(true); + when(buffer.size()).thenReturn(0); + when(buffer.nextBatch(anyBoolean())).thenReturn(null); + } else { + buffer = + new BatchedLogBuffer( + maxBatchSize, + maxBytes, + maxAge, + queueCapacity, + metrics, + mockBufferClock ? mockClock : BatchedLogBuffer.DEFAULT_CLOCK); + } + var uploader = mock(AsyncBatchedLogUploader.class); + var logHandler = + new BillingUploadingLogHandler( + buffer, + uploader, + shortUploadSleepDuration ? UPLOAD_SLEEP_DURATION_SHORT : UPLOAD_SLEEP_DURATION, + UPLOADER_SAFETY_DEADLINE, + shortUploadShutdownDeadline + ? UPLOAD_SHUTDOWN_DEADLINE_SHORT + : UPLOAD_SHUTDOWN_DEADLINE); + + return new Fixture( + maxBatchSize, + maxBytes, + maxAge, + queueCapacity, + logRecords, + buffer, + metrics, + logHandler, + uploader, + mockClock); + } + + /** Assert the buffer is full, and so offer() fails */ + void assertBufferFull(String desc, int index) { + + // although the next log record is wafer-thin, it is too much for Mr Creosote + assertThat(buffer().offer(logRecords.get(index))).as(desc + " - fail at capacity").isFalse(); + + // Running again to confirm it is still full + assertThat(buffer().offer(logRecords.get(index))) + .as(desc + " - second - fail at capacity") + .isFalse(); + } + + /** + * Offer the log records selected by slice to the buffer, all should work, assert the buffer has + * the items the slice selected + */ + void assertOffer(String desc, BatchedLogBufferTest.Slice slice) { + + var snapshot = BufferSnapshot.create(this); + + for (var record : slice.stream(logRecords).toList()) { + assertThat(buffer.offer(record)).as(desc + " - assertOffer() - offering").isTrue(); + } + + snapshot.assertAll(desc, slice, true); + } + + /** + * Get a batch from the buffer, assert we got a batch that is legal, and assert the buffer has + * changed by the amount of the batch + */ + BatchedLogBuffer.Batch assertNextBatch(String desc, boolean drainFully) { + + var snapshot = BufferSnapshot.create(this); + var batch = buffer.nextBatch(drainFully); + + // assert the batch is what we expected. + assertThat(batch).as(desc + " - assertNextBatch() - batch is not null").isNotNull(); + + assertThat(batch.size()) + .as(desc + " - assertNextBatch() - batch size <= MAX_BATCH_SIZE") + .isLessThanOrEqualTo(maxBatchSize); + // note: it is legal to have a batch bigger than the maxBytes, specialised tests for that + // shoudl only happen when there is a single log record bigger than maxBytes + assertThat(batch.bytes()) + .as(desc + " - assertNextBatch() - batch bytes <= MAX_BATCH_BYTES") + .isLessThanOrEqualTo(maxBytes); + + // assert the buffer updated bookkeeping as we expect + snapshot.assertAll(desc, batch); + return batch; + } + + // just redeclare with no exception to make it easier + interface NoExceptionCloseable extends AutoCloseable { + @Override + void close(); + } + + /** + * Start the logHandler uploading on a daemon thread, and returns a closeable for killing the + * thread. + * + * @return + */ + NoExceptionCloseable startHandlerUploading(String desc) { + + var executor = + Executors.newSingleThreadExecutor(Thread.ofPlatform().daemon().name(desc).factory()); + + var uploadingFuture = + executor.submit( + () -> { + LOGGER.info("startHandlerUploading() - starting logHandler. desc:{}", desc); + logHandler.startUploading(); + }); + + return () -> { + LOGGER.info("startHandlerUploading() - stopping logHandler. desc:{}", desc); + try { + // this is waiting for the uploading thread to return + uploadingFuture.get(10, TimeUnit.SECONDS); + executor.awaitTermination(1000, TimeUnit.MILLISECONDS); + + LOGGER.info("startHandlerUploading() - stopped logHandler. desc:{}", desc); + + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + LOGGER.error("startHandlerUploading() - error 1", e); + throw new RuntimeException(e); + } catch (ExecutionException e) { + LOGGER.error("startHandlerUploading() - error 2", e); + + throw e.getCause() instanceof RuntimeException re + ? re + : new RuntimeException(e.getCause()); + } catch (TimeoutException e) { + LOGGER.error("startHandlerUploading() - error 3", e); + + throw new RuntimeException(e); + } finally { + executor.shutdownNow(); + } + }; + } + } + + /** + * A slice of a list, `from` is inclusive, `to` is exclusive + * + *

... + */ + record Slice(int from, int to) { + + public Stream stream(List list) { + return list.stream().skip(from).limit(to - from); + } + + public int size() { + return to - from; + } + + public static Slice to(int to) { + return new Slice(0, to); + } + + public static Slice from(int from) { + return new Slice(from, Integer.MAX_VALUE); + } + + public static Slice slice(int from, int to) { + return new Slice(from, to); + } + } + + /** + * Snapshot of the metadata (size etc) for the buffer, that can be used to compare how the buffer + * metadata has changed + * + *

... + */ + record BufferSnapshot( + boolean isEmpty, int size, long queuedBytes, int remainingCapacity, Fixture fixture) { + + static BufferSnapshot create(Fixture fixture) { + // reset the counters for calls to metrics + clearInvocations(fixture.metrics); + return new BufferSnapshot( + fixture.buffer().isEmpty(), + fixture.buffer().size(), + fixture.buffer().queuedBytes(), + fixture.buffer().remainingCapacity(), + fixture); + } + + /** + * Assert that the current metadata values for the buffer are the values in the snapshot PLUS + * the log records that were added by the Slice. + */ + void assertAll(String desc, Slice slice, boolean inOrder) { + assertBufferMetadata(desc, slice); + assertBufferItems(desc, slice, inOrder); + } + + /** + * Assert that the current metadata values for the buffer are the values in the snapshot MINUS + * the buffer entries that were removed in the batch + */ + void assertAll(String desc, BatchedLogBuffer.Batch batch) { + assertBufferMetadata(desc, batch); + assertBufferItems(desc, batch); + } + + /** current buffer metadata = snapshot + slice */ + void assertBufferMetadata(String desc, Slice slice) { + + if (slice.size() == 0) { + assertThat(fixture.buffer().isEmpty()) + .as(desc + " - isEmpty no change after empty slice") + .isEqualTo(isEmpty()); + } else { + assertThat(fixture.buffer().isEmpty()) + .as(desc + " - isEmpty false after non empty slice") + .isEqualTo(false); + } + + assertThat(fixture.buffer().size()) + .as(desc + " - post buffer size increased by slice") + .isEqualTo(size() + slice.size()); + + verify( + fixture.metrics, + times(slice.size()).description(desc + "metrics called for every offer")) + .offered(); + + long addedBytes = 0; + for (var record : slice.stream(fixture.logRecords).toList()) { + addedBytes += BatchedLogBuffer.Entry.lineBytes(record.getMessage()); + } + + assertThat(fixture.buffer().queuedBytes()) + .as(desc + " - post buffer bytes increased by slice") + .isEqualTo(queuedBytes + addedBytes); + } + + /** current buffer metadata = snapshot - batch */ + void assertBufferMetadata(String desc, BatchedLogBuffer.Batch batch) { + + assertThat(fixture.buffer().size()) + .as(desc + " - buffer size decreased by batch size") + .isEqualTo(size() - batch.size()); + + assertThat(fixture.buffer().queuedBytes()) + .as(desc + " - buffer bytes size decreased by batch bytes") + .isEqualTo(queuedBytes - batch.bytes()); + } + + /** + * current buffer items contain items from slice inOrder - if we expect items in buffer to match + * order of the fixture + */ + void assertBufferItems(String desc, Slice slice, boolean inOrder) { + + var bufferItems = fixture.buffer().peekBuffer(); + + int i = slice.from() > bufferItems.size() ? 0 : slice.from(); + for (var record : slice.stream(fixture.logRecords).toList()) { + + if (inOrder) { + assertThat(record.getMessage()) + .as(desc + " - buffer items at position match exactly pos: " + i) + .isEqualTo(bufferItems.get(i++).line()); + } else { + + var entry = new BatchedLogBuffer.Entry(record.getInstant(), record.getMessage()); + assertThat(bufferItems) + .as(desc + " - buffer items contains entry: " + entry) + .contains(entry); + } + } + } + + /** current buffer items contain NONE of items in batch */ + void assertBufferItems(String desc, BatchedLogBuffer.Batch batch) { + + var peekedBuffer = fixture.buffer().peekBuffer(); + + for (var batchString : batch.lines()) { + + var found = peekedBuffer.stream().anyMatch(entry -> entry.line().equals(batchString)); + assertThat(found) + .as(desc + " - line from batch no longer in buffer: " + batchString) + .isFalse(); + } + } + } +} diff --git a/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingUploadingLogHandlerTest.java b/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingUploadingLogHandlerTest.java new file mode 100644 index 0000000000..7113c62967 --- /dev/null +++ b/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingUploadingLogHandlerTest.java @@ -0,0 +1,309 @@ +package io.stargate.sgv2.jsonapi.service.billing; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; +import static org.mockito.Mockito.*; + +import io.smallrye.mutiny.Uni; +import java.util.List; +import java.util.logging.LogRecord; +import java.util.stream.IntStream; +import org.junit.jupiter.api.Test; +import org.mockito.ArgumentCaptor; +import org.mockito.verification.VerificationMode; + +/** */ +public class BillingUploadingLogHandlerTest extends BillingTestBase { + + // ********************************************************* + // Handler interface - Producer side of the handler + // ********************************************************* + + /** Null record silently dropped by handler */ + @Test + public void publishSilentDropNulls() { + + var fixture = defaultLogHandlerFixture(); + + fixture.logHandler().publish(null); + fixture.logHandler().publish(null); + fixture.logHandler().publish(null); + + verify( + fixture.buffer(), + times(0).description("publishSilentDropNulls() - no calls to buffer.offer()")) + .offer(any()); + } + + /** Non null record silently dropped by handler when closed */ + @Test + public void publishSilentDropWhenClosed() { + + var fixture = defaultLogHandlerFixture(); + fixture.logHandler().close(); + + fixture.logHandler().publish(fixture.logRecords().getFirst()); + fixture.logHandler().publish(fixture.logRecords().getFirst()); + fixture.logHandler().publish(fixture.logRecords().getFirst()); + + verify( + fixture.buffer(), + times(0).description("publishSilentDropWhenClosed() - no calls to buffer.offer()")) + .offer(any()); + } + + /** Records published when buffer is full are dropped, no error */ + @Test + public void publishSilentWhenBufferFull() { + + var fixture = defaultLogHandlerFixture(); + + // return false, buffer full , go away + when(fixture.buffer().offer(any())).thenReturn(false); + + fixture.logHandler().publish(fixture.logRecords().getFirst()); + fixture.logHandler().publish(fixture.logRecords().getFirst()); + fixture.logHandler().publish(fixture.logRecords().getFirst()); + + verify( + fixture.buffer(), + times(3).description("publishSilentWhenBufferFull() - called offer for each record")) + .offer(fixture.logRecords().getFirst()); + } + + /** Passing record to handler, is then passed to the buffer. */ + @Test + public void publishOfferSucceed() { + + var fixture = defaultLogHandlerFixture(); + var slice = Slice.to(MAX_BATCH_BYTES_NUM_MESSAGES); + var expectedLogRecords = slice.stream(fixture.logRecords()).toList(); + + var argCaptor = ArgumentCaptor.forClass(LogRecord.class); + + for (var record : expectedLogRecords) { + fixture.logHandler().publish(record); + } + + verify( + fixture.buffer(), + times(expectedLogRecords.size()) + .description("publishOfferSucceed() - buffer called for each log record")) + .offer(argCaptor.capture()); + var actualLogRecords = argCaptor.getAllValues(); + + assertThat(actualLogRecords) + .as("publishOfferSucceed() - all and only expected records passed to the buffer") + .containsExactlyElementsOf(expectedLogRecords); + } + + /** Calling flush on handler that is NOT uploading does nothing */ + @Test + public void flushNullOpIfNotStarted() { + + var fixture = defaultLogHandlerFixture(); + + fixture.logHandler().flush(); + // there is no uploading, so should not ask for next batch + verify(fixture.buffer(), never()).nextBatch(anyBoolean()); + } + + /** + * Verify the number of times a function was called, but with a timeout to wait. e.g. when waiting + * for the uploading thread to wakeup + */ + private VerificationMode timeoutTimes(String desc, int times) { + return timeout(2000).times(times).description(desc); + } + + /** Calling flush on handler that is uploading causes handler to check buffer. */ + @Test + public void flushChecksForBatch() { + + var fixture1 = defaultLogHandlerFixture(); + try (var handlerThread = + fixture1.startHandlerUploading("flushChecksForBatch() - close not called")) { + fixture1.logHandler().flush(); + // close has not been called, so it should not drain + verify(fixture1.buffer(), timeoutTimes("nextBatch() called once with drain false", 1)) + .nextBatch(false); + + // this is a bit stupid, calling close to close the upload thread + fixture1.logHandler().close(); + } + + var fixture2 = defaultLogHandlerFixture(); + try (var handlerThread = + fixture2.startHandlerUploading("flushChecksForBatch() - close is called")) { + fixture2.logHandler().unsafeClose(); + fixture2.logHandler().flush(); + // close has been called, so it should drain buffer + verify(fixture2.buffer(), timeoutTimes("nextBatch() called once with drain true", 1)) + .nextBatch(true); + } + } + + @Test + public void closeWithoutUploadThreadReturns() { + + var fixture1 = defaultLogHandlerFixture(); + // NOT STARTING upload + fixture1.logHandler().close(); + // upload not running, should not try to get a batch + verify(fixture1.buffer(), timeoutTimes("nextBatch() never called", 0)).nextBatch(anyBoolean()); + // should have closed the uploader + verify(fixture1.uploader(), timeoutTimes("uploader.close() called", 1)).close(); + } + + /** + * Call close, but the upload thread has not released the upload permit, so close cannot detect + * upload has finished. + */ + @Test + public void closeReturnsWhenUploadUnstopped() { + + var fixture1 = defaultLogHandlerFixture(true, false, true); + // NOT STARTING upload, but acquire the permit it would take + fixture1.logHandler().unsafeAcquireUploadPermit(); + // close() will not return until it times out waiting for the upload thread to finish + fixture1.logHandler().close(); + } + + /** + * Calling close() when the handler is running should cause the buffer to be called to drain it. + */ + @Test + public void closeCausesBufferDrain() { + + var fixture1 = defaultLogHandlerFixture(); + try (var handlerThread = fixture1.startHandlerUploading("closeCausesBufferDrain()")) { + threadSleep(100); // give the uploader time to get into the wait on wakeup + + fixture1.logHandler().close(); + // close has been called, so it should drain + verify(fixture1.buffer(), timeoutTimes("nextBatch() called with drain=true", 1)) + .nextBatch(true); + verify(fixture1.uploader(), timeoutTimes("uploader.close() called", 1)).close(); + } + // the auto closable will wait for the upload thread to naturally exit + } + + // ********************************************************* + // startUploading - Consumer side of the handler + // ********************************************************* + + /** + * Calling startUpLoad twice on different threads, fails because there can be only one active + * thread running the function + * + *

Cannot call on same thread as it will be parked running the upload + */ + @Test + public void startUploadingCalledTwiceFails() { + + var fixture1 = defaultLogHandlerFixture(); + try (var handlerThread1 = + fixture1.startHandlerUploading("startUploadingCalledTwiceFails() - 1st")) { + // make sure the worker thread has time to start + threadSleep(10); + + // the exception will happen when startUploading is entered, but we + // wont get the error until calling close() which calls Future.get() + var closable = fixture1.startHandlerUploading("startUploadingCalledTwiceFails() - 2nd"); + // make sure the worker thread has time to start + threadSleep(10); + + assertThatThrownBy(closable::close, "startUploadingCalledTwiceFails() - second call") + .isInstanceOf(IllegalStateException.class); + + // stop the first thread that is running startUpload() + fixture1.logHandler().close(); + } + } + + /** In normal operation startUpload detects three batches and sends to uploader */ + @Test + public void startUploadingSendsToUploader() { + + var NUM_BATCHES = 3; + var fixture1 = defaultLogHandlerFixture(); + try (var handlerThread1 = + fixture1.startHandlerUploading("startUploadingSendsToUploader() - upload thread")) { + // make sure the worker thread has time to start and get to the sleep. + threadSleep(100); + + // ** TESTING NORMAL OPERATION + + // setup buffer to return three batches we will collect in normal operations + var expectedNormalBatches = mockUploading(fixture1, NUM_BATCHES, false); + // handler should be sleeping because of long sleep, flush will wake it up. + fixture1.logHandler().flush(); + // wait for it the handler to call the uploader + var normalCaptor = ArgumentCaptor.forClass(BatchedLogBuffer.Batch.class); + verify( + fixture1.uploader(), + timeoutTimes( + "startUploadingSendsToUploader() - normal mode", expectedNormalBatches.size())) + .upload(normalCaptor.capture()); + var actualNormalBatches = normalCaptor.getAllValues(); + + // ** TESTING CLOSE / SHUTDOWN OPERATION + + // reset counter , will already be NUM_BATCHES from the normal operation check above + clearInvocations(fixture1.uploader()); + // we are using the long upload sleep, uploader should be asleep again, send batches + // and close to see we get correct behavior + var expectedShutdownBatches = mockUploading(fixture1, NUM_BATCHES, true); + // close to get shutdown operations + fixture1.logHandler().close(); + // wait for it the handler to call the uploader + var shutdownCaptor = ArgumentCaptor.forClass(BatchedLogBuffer.Batch.class); + verify( + fixture1.uploader(), + timeoutTimes( + "startUploadingSendsToUploader() - shutdown mode", + expectedShutdownBatches.size())) + .upload(shutdownCaptor.capture()); + var actualShutdownBatches = shutdownCaptor.getAllValues(); + + assertThat(actualNormalBatches) + .as("startUploadingSendsToUploader() - batches from normal operation match") + .containsExactlyElementsOf(expectedNormalBatches); + + assertThat(actualShutdownBatches) + .as("startUploadingSendsToUploader() - batches from shutdown operation match") + .containsExactlyElementsOf(expectedShutdownBatches); + + // we have closed handler, should exit block now + } + } + + private List mockUploading( + Fixture fixture, int numBatches, boolean expectDrainFully) { + // by default the mock buffer will be returning null for nextBatch(), the + // log handler should be in a wait because we are using long upload sleep + + var expectedBatches = + IntStream.range(0, numBatches).mapToObj(i -> mock(BatchedLogBuffer.Batch.class)).toList(); + + // we expect the uploader to be called with these batches, and it needs to + // return a Uni with an UploadResult - but we dont need to keep the UploadResult + // MUST set up the uploader to return a value before putting batches into the buffer + for (var batch : expectedBatches) { + var result = new AsyncBatchedLogUploader.UploadResult(batch, null); + when(fixture.uploader().upload(batch)).thenReturn(Uni.createFrom().item(result)); + } + + // now connect the Batch's to the buffer so it will return them. Once this is done + // the upload thread will pick them up once woken + // expectDrainFully depends on if the test is closed the handler. + var nextBatchStub = when(fixture.buffer().nextBatch(expectDrainFully)); + for (var batch : expectedBatches) { + nextBatchStub = nextBatchStub.thenReturn(batch); + } + // now return null as sentinel for no more data + nextBatchStub.thenReturn(null); + + return expectedBatches; + } +}