NOTE: This enables the collection of billing events and configures their structure, the + * events are then send to the logger named {@link + * io.stargate.sgv2.jsonapi.service.billing.DefaultBilling#BILLING_LOGGER_NAME}. Use {@link + * BillingS3UploadConfig} to configure sending them to S3, or configure a log handler manually. + */ @ConfigMapping(prefix = "stargate.jsonapi.billing") public interface BillingConfig { diff --git a/src/main/java/io/stargate/sgv2/jsonapi/config/BillingS3UploadConfig.java b/src/main/java/io/stargate/sgv2/jsonapi/config/BillingS3UploadConfig.java new file mode 100644 index 0000000000..7bedf55f3e --- /dev/null +++ b/src/main/java/io/stargate/sgv2/jsonapi/config/BillingS3UploadConfig.java @@ -0,0 +1,122 @@ +package io.stargate.sgv2.jsonapi.config; + +import io.smallrye.config.ConfigMapping; +import io.smallrye.config.WithDefault; +import io.stargate.sgv2.jsonapi.service.billing.BillingS3Lifecycle; +import io.stargate.sgv2.jsonapi.service.billing.BillingUploadingLogHandler; +import java.time.Duration; +import java.util.Optional; +import software.amazon.awssdk.core.retry.RetryMode; + +/** + * Configuration for the billing S3 export (see {@link BillingS3Lifecycle}). + * + *
See also {@link BillingConfig}
+ */
+@ConfigMapping(prefix = "stargate.jsonapi.billing.s3")
+public interface BillingS3UploadConfig {
+
+ /** When NOTE: A single instance of the buffer is expected to as the application runs, and it
+ * must call {@link #registerBuffer(BatchedLogBuffer)} when created.
+ *
+ * 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 final class BatchedLogBufferMetrics extends MetricsBase {
+
+ private final AtomicBoolean bufferRegister = new AtomicBoolean(false);
+
+ public final Counter offered;
+ public final Counter dropped;
+
+ // these are nullable, they are set when registerBuffer is called.
+ public TimeGauge headAgeMs;
+ public Gauge size;
+ public Gauge remainingCapacity;
+ public Gauge bytes;
+
+ /** Creates a new instance, remember to call {@link #registerBuffer(BatchedLogBuffer)} */
+ public BatchedLogBufferMetrics(MeterRegistry meterRegistry, String prefix) {
+ super(meterRegistry, prefix);
+
+ this.offered = newCounter("buffer.offered");
+ this.dropped = newCounter("buffer.dropped");
+ }
+
+ /**
+ * Registers the supplued buffer as the source for metrics thats are pulled from the buffer. Can
+ * only be called once.
+ */
+ public void registerBuffer(BatchedLogBuffer buffer) {
+
+ if (!bufferRegister.compareAndSet(false, true)) {
+ throw new IllegalStateException("registerBuffer() already called");
+ }
+
+ headAgeMs =
+ newTimeGauge(
+ "buffer.head_age_ms", () -> buffer.headEntryAge().toMillis(), TimeUnit.MILLISECONDS);
+ size = newGauge("buffer.size", buffer::size);
+ remainingCapacity = newGauge("buffer.remaining_capacity", buffer::remainingCapacity);
+ bytes = newGauge("buffer.bytes", buffer::queuedBytes);
+ }
+
+ public void incrementOffered() {
+ offered.increment();
+ }
+
+ public void incrementDropped() {
+ 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..c7141123d0
--- /dev/null
+++ b/src/main/java/io/stargate/sgv2/jsonapi/metrics/BatchedLogUploaderMetrics.java
@@ -0,0 +1,50 @@
+package io.stargate.sgv2.jsonapi.metrics;
+
+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 uploading billing events, mostly around what is sent to S3 See {@link
+ * io.stargate.sgv2.jsonapi.service.billing.S3BatchedLogUploader}
+ */
+public final class BatchedLogUploaderMetrics extends MetricsBase {
+
+ public final Counter uploadedBatches;
+ public final Counter uploadedBytes;
+ public final Timer uploadedHeadAgeMs;
+ public final Counter uploadedEvents;
+
+ public final Counter failedBatches;
+ public final Counter failedEvents;
+
+ public BatchedLogUploaderMetrics(MeterRegistry meterRegistry, String prefix) {
+ super(meterRegistry, prefix);
+
+ this.uploadedBatches = newCounter("uploaded.batches");
+ this.uploadedBytes = newCounter("uploaded.bytes");
+ this.uploadedEvents = newCounter("uploaded.events");
+ this.uploadedHeadAgeMs = newTimer("uploaded.oldest_event");
+
+ this.failedBatches = newCounter("failed.batches");
+ this.failedEvents = newCounter("failed.events");
+ }
+
+ public void recordBatchDelivered(BatchedLogBuffer.Batch batch) {
+
+ Objects.requireNonNull(batch, "batch must not be null");
+ uploadedBatches.increment();
+ uploadedBytes.increment(batch.bytes());
+ uploadedHeadAgeMs.record(batch.oldestEventAtDuration());
+ uploadedEvents.increment(batch.size());
+ }
+
+ 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..3882585a42
--- /dev/null
+++ b/src/main/java/io/stargate/sgv2/jsonapi/metrics/MetricsBase.java
@@ -0,0 +1,72 @@
+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
+ *
+ * Use this so that metrics have a common name, and use the factory functions when creating new
+ * metrics so we can push them through a single point of control
+ */
+public abstract class MetricsBase {
+
+ protected final MeterRegistry meterRegistry;
+ protected final String prefix;
+
+ /**
+ * Creates a new instance.
+ *
+ * @param meterRegistry Registry that metrics will be registered with.
+ * @param prefix Prefix that is given to all metrics created with the factory functions
+ */
+ 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 There are two users of this class:
+ *
+ * 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 with configuration.
+ *
+ * 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 more than maxBatchBytes if there is a single log record that is bigger.
+ * @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.
+ * @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 PROD CODE. If null uses {@link
+ * #DEFAULT_CLOCK}
+ */
+ @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);
+ }
+ if (maxBatchBytes < 1) {
+ throw new IllegalArgumentException("maxBatchBytes must be >= 1, got: " + maxBatchBytes);
+ }
+ if (maxBatchAge == null || maxBatchAge.isNegative() || maxBatchAge.isZero()) {
+ throw new IllegalArgumentException("maxAge must be positive, got: " + maxBatchAge);
+ }
+
+ this.maxBatchSize = maxBatchSize;
+ this.maxBatchBytes = maxBatchBytes;
+ this.maxBatchAge = maxBatchAge;
+ this.capacity = capacity;
+ this.metrics = Objects.requireNonNull(metrics, "billingMetrics must not be null");
+
+ 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<>(capacity);
+
+ // just to be safe, register after queue created in case metrics are scrapped
+ this.metrics.registerBuffer(this);
+ }
+
+ /**
+ * 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.
+ */
+ 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);
+
+ metrics.incrementOffered();
+ if (!queue.offer(newEntry)) {
+ // Bounded buffer full, drop and count
+ LOGGER.warn("offer() - buffer full, dropping new entry: {}", newEntry);
+ metrics.incrementDropped();
+ 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 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
+ * 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() {
+ return entryAge(queue.peek());
+ }
+
+ @VisibleForTesting
+ Duration entryAge(Entry entry) {
+ return entry == null ? Duration.ZERO : Duration.between(entry.eventAt(), clock.instant()).abs();
+ }
+
+ private BillingBatchReason decideNextBatch(boolean drainFully) {
+
+ 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 (LOGGER.isTraceEnabled()) {
+ LOGGER.trace("decideNextBatch() - drainFully:{} , decision:{}", drainFully, decision);
+ }
+ 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_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 id = UUID_V7_GENERATOR.generate();
+ private final BillingBatchReason reason;
+ private final List 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 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/provider/Billing.java b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/Billing.java
similarity index 55%
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..6ae650596a 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,35 +1,22 @@
-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;
/**
* Per-request billing sink for model calls.
*
* Get an instance via {@code requestContext.billing()}, which delegates to {@link
- * #create(BillingConfig, ApiFeatures)} to pick the right implementation for the request:
+ * #create(BillingConfig, ApiFeatures)} to pick the right implementation for the request
*
- * Call {@link #emitEvent(ModelUsage)} to generate a billing event for the model usage, the
+ * implementation takes care of where the event goes.
*/
public interface Billing {
- /**
- * Emits billing events for the given aggregated model usage.
- *
- * @param modelUsage usage data for the model call; must not be null.
- */
- void emitEvent(ModelUsage modelUsage);
-
/**
* Shared NO-OP {@link Billing}. Still enforces the non-null {@code modelUsage} contract so tests
* (and the feature-disabled production path) don't accidentally mask null-passing bugs in calling
@@ -38,17 +25,29 @@ public interface Billing {
Billing NO_OP = modelUsage -> Objects.requireNonNull(modelUsage, "modelUsage must not be null");
/**
- * Factory that picks the right {@link Billing} implementation for the current request.
- * Centralizes the {@code DefaultBilling vs NO_OP} dispatch so callers (e.g. {@link
- * io.stargate.sgv2.jsonapi.api.request.RequestContext}) don't have to know the rule. Reads {@code
- * config} only when the feature is enabled — when disabled it is fine to pass any value,
- * including one that would not validate as a real config.
+ * Emits billing events for the given model usage.
+ *
+ * @param modelUsage usage data for the model call; must not be null.
+ */
+ void emitEvent(ModelUsage modelUsage);
+
+ /**
+ * Factory to create the correct instance of {@link Billing} for a given request.
+ *
+ * Each event represents a single billable unit from a user request (e.g. token usage or egress
- * bytes for an embedding or reranking call). Constructed and emitted by {@link Billing}.
+ * bytes for an embedding or reranking call). Constructed and emitted by {@link DefaultBilling}.
*
* Example serialized form:
*
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 94%
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..9626f46b62 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;
@@ -43,8 +43,7 @@ public enum Metric {
private final Metric metric;
BillingEventType(String eventName, boolean internal, Metric metric) {
- // Event names are emitted lower-case in the JSON billing event payload; normalize defensively
- // so a stray uppercase character in the literal can't break downstream consumers.
+ // Event names are emitted lower-case in the JSON billing event payload
this.eventName = eventName.toLowerCase();
this.internal = internal;
this.metric = metric;
diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3Lifecycle.java b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3Lifecycle.java
new file mode 100644
index 0000000000..d07ff50d92
--- /dev/null
+++ b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3Lifecycle.java
@@ -0,0 +1,142 @@
+package io.stargate.sgv2.jsonapi.service.billing;
+
+import com.google.common.annotations.VisibleForTesting;
+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.BillingS3UploadConfig;
+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;
+import java.util.logging.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Lifecycle management for sending billing events to S3.
+ *
+ * If configured to upload to S3, a long-running task is started using the {@link
+ * Infrastructure#getDefaultWorkerPool()} in the background for the {@link
+ * BillingUploadingLogHandler#startUploading()} to run on.
+ *
+ * S3 uploading is controlled with the {@link BillingS3UploadConfig}
+ *
+ * See {@link BillingS3Lifecycle} for setup.
+ *
+ * Designed to be used in two ways:
+ *
+ * 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 the uploader.
+ */
+ @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
+ 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.
+ *
+ * Call {@link #close()} to signal this function to stop waiting, upload all events, and
+ * return.
+ */
+ void startUploading() {
+
+ LOGGER.info("startUploading() - handler:{}, buffer:{}, uploader:{}", this, buffer, uploader);
+
+ boolean hasUploadPermit = false;
+ BatchedLogBuffer.Batch batch;
+ try {
+ if (!(hasUploadPermit = 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 {
+ 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, to signal close() to stop waiting
+ if (hasUploadPermit) {
+ 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. handler:{}, buffer:{}, uploader:{}",
+ this,
+ buffer,
+ uploader);
+ }
+
+ /** Signals the thread running {@link #startUploading()} to wake up and check for work. */
+ private void notifyUploading() {
+ // Adds a permit to the wakeupPermit so the uploading thread will wakeup and do some work.
+ 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. */
+ 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 in case it breaks.
+ // we do not want the uploader hanging to block this thread.
+
+ // 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, this causes deferred to start exec'ing
+
+ 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());
+ dumpBuffer();
+ }
+
+ /** TODO: dump buffer or a failed batch to regular logs or whatever */
+ private void dumpBuffer() {
+ LOGGER.warn("dumpBuffer() - DUMP BUFFER CALLED, NOT IMPLEMENTED YET!");
+ }
+
+ private void dumpBatch(BatchedLogBuffer.Batch batch) {
+ LOGGER.warn("dumpBuffer() - DUMP BATCH CALLED, NOT IMPLEMENTED YET!");
+ }
+
+ @Override
+ public String toString() {
+ return new StringBuilder(classSimpleName(this) + "{")
+ .append("uploadSleepDuration=")
+ .append(uploadSleepDuration)
+ .append(", uploadShutdownDeadline=")
+ .append(uploadShutdownDeadline)
+ .append(", uploaderSafetyDeadline=")
+ .append(uploaderSafetyDeadline)
+ .append(", isClosed=")
+ .append(isClosed)
+ .append(", wakeupPermit.availablePermits=")
+ .append(wakeupPermit.availablePermits())
+ .append(", uploadPermit.availablePermits=")
+ .append(uploadPermit.availablePermits())
+ .append("}")
+ .toString();
+ }
+}
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
new file mode 100644
index 0000000000..2b58d8401e
--- /dev/null
+++ b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/DefaultBilling.java
@@ -0,0 +1,131 @@
+package io.stargate.sgv2.jsonapi.service.billing;
+
+import static io.stargate.sgv2.jsonapi.util.StringUtil.requireNonBlank;
+
+import com.fasterxml.jackson.core.JacksonException;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.databind.ObjectWriter;
+import com.google.common.annotations.VisibleForTesting;
+import io.stargate.sgv2.jsonapi.config.BillingConfig;
+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 java.util.function.Function;
+import org.apache.commons.lang3.tuple.Pair;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * {@link Billing} implementation that emits structured JSON log lines on the {@code billing.events}
+ * logger for downstream billing pipelines.
+ *
+ * For each {@link ModelUsage}, up to three events are emitted, for the egress and ingress bytes
+ * of the request and the total number of tokens used. See {@link BillingEventType} for which ones
+ * should be used.
+ *
+ * NOTE: this class relies on the configured format for the billing logger called {@link
+ * #BILLING_LOGGER_NAME}, it should only log the message on the line because that has the JSON of
+ * the {@link BillingEvent}
+ */
+public class DefaultBilling implements Billing {
+
+ /** Name of the logger that billing events will be sent to. */
+ public static final String BILLING_LOGGER_NAME = "billing.events";
+
+ private static final Logger BILLING_LOGGER = LoggerFactory.getLogger(BILLING_LOGGER_NAME);
+ private static final Logger LOGGER = LoggerFactory.getLogger(DefaultBilling.class);
+
+ private static final ObjectWriter OBJECT_WRITER = new ObjectMapper().writer();
+
+ // There are three metrics, but 6 Event Types, because 3 each for internal or external provider.
+ // this is the mapping of the metric and where we get it from
+ private static final List All events from a single {@code ModelUsage} share one timestamp so they can be correlated in
+ * the billing logs.
+ */
+ @VisibleForTesting
+ List TODO: checksums ? .requestChecksumCalculation(RequestChecksumCalculation.WHEN_SUPPORTED)
+ * .responseChecksumValidation(ResponseChecksumValidation.WHEN_SUPPORTED)
+ */
+public class S3BatchedLogUploader implements AsyncBatchedLogUploader {
+
+ private static final Logger LOGGER = LoggerFactory.getLogger(S3BatchedLogUploader.class);
+
+ // this header not defined in standard libs
+ @VisibleForTesting public static final String CONTENT_TYPE_NDJSON = "application/x-ndjson";
+ private static final DateTimeFormatter OBJECT_KEY_FORMATTER =
+ DateTimeFormatter.ofPattern("yyyy/MM/dd/HH/mm").withZone(ZoneOffset.UTC);
+
+ private final S3AsyncClient client;
+ private final String region;
+ private final String bucket;
+ private final String pathPrefix;
+ private final BatchedLogUploaderMetrics metrics;
+
+ /**
+ * Visible for testing, use the {@link #create(String, String, String, String, Duration, Duration,
+ * RetryMode, BatchedLogUploaderMetrics)} to get a new instance.
+ */
+ @VisibleForTesting
+ protected S3BatchedLogUploader(
+ S3AsyncClient client,
+ String region,
+ String bucket,
+ String pathPrefix,
+ BatchedLogUploaderMetrics metrics) {
+ this.client = client;
+ this.region = region;
+ this.bucket = bucket;
+ this.pathPrefix = pathPrefix;
+ this.metrics = metrics;
+ }
+
+ /**
+ * Factory to create a new instance, including creating and configuring the AWS client
+ *
+ * @param region S3 region to send events to
+ * @param bucket S3 bucket to send events to
+ * @param endpointOverride nullable override for the S3 end point, used for testing.
+ * @param pathPrefix Prefix to store events under, the path also has the event time and a UUId in
+ * it.
+ * @param s3CallAttemptTimeout Timeout for a single call to S3
+ * @param s3TotalCallTimeout Timeout for the total time calling s3, including retries
+ * @param s3RetryMode retry strategy, see S3 docs.
+ * @param metrics Metrics instance to track with.
+ * @return
+ */
+ public static S3BatchedLogUploader create(
+ String region,
+ String bucket,
+ String endpointOverride,
+ String pathPrefix,
+ Duration s3CallAttemptTimeout,
+ Duration s3TotalCallTimeout,
+ RetryMode s3RetryMode,
+ BatchedLogUploaderMetrics metrics) {
+
+ if (region == null || region.isBlank()) {
+ throw new IllegalArgumentException("region must not be null or blank");
+ }
+ if (bucket == null || bucket.isBlank()) {
+ throw new IllegalArgumentException("bucket must not be null or blank");
+ }
+ if (pathPrefix == null || pathPrefix.isBlank()) {
+ throw new IllegalArgumentException("pathPrefix must not be null or blank");
+ }
+ Objects.requireNonNull(s3TotalCallTimeout, "s3TotalCallTimeout must not be null");
+ Objects.requireNonNull(s3CallAttemptTimeout, "s3CallAttemptTimeout must not be null");
+ Objects.requireNonNull(s3RetryMode, "s3RetryMode 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().
+ LOGGER.info(
+ "create() - region:{}, bucket:{}, pathPrefix:{}, s3TotalCallTimeout:{}, s3CallAttemptTimeout:{}, s3RetryMode:{}",
+ region,
+ bucket,
+ pathPrefix,
+ s3TotalCallTimeout,
+ s3CallAttemptTimeout,
+ s3RetryMode);
+
+ var builder =
+ S3AsyncClient.builder()
+ .region(Region.of(region))
+ .overrideConfiguration(
+ config ->
+ config
+ .retryStrategy(s3RetryMode)
+ .apiCallAttemptTimeout(s3CallAttemptTimeout)
+ .apiCallTimeout(s3TotalCallTimeout));
+
+ // Real AWS S3 needs no endpoint: the SDK endpoint rules (s3 SDK's DefaultS3EndpointProvider)
+ // derive https://
+ *
+ * @param batch The batch of log records to upload
+ * @return Uni for uploading the batch to S3. The returned Uni will be working to complete the
+ * task (i.e. not deferred), and will be emitted on the {@link
+ * Infrastructure#getDefaultWorkerPool()}
+ */
+ @Override
+ public Uni Construction is driven by {@link
- * io.stargate.sgv2.jsonapi.api.request.RequestContext#billing()}: the request context decides
- * between this implementation and {@link Billing#NO_OP} based on whether {@link
- * ApiFeature#BILLING_EVENTS_LOGGING} is enabled for the request. This class therefore assumes
- * billing is enabled and unconditionally emits — it does not re-check the feature flag.
- *
- * For each {@link ModelUsage}, up to three events are emitted, one per billable metric ({@link
- * BillingEventType.Metric#TOTAL_TOKENS TOTAL_TOKENS}, {@link BillingEventType.Metric#EGRESS_BYTES
- * EGRESS_BYTES}, {@link BillingEventType.Metric#INGRESS_BYTES INGRESS_BYTES}). The {@code
- * internal_*} variant is used when the model provider is listed in {@link
- * BillingConfig#internalModelProviders()}; otherwise the {@code external_*} variant is used. Events
- * whose type is not in {@link BillingConfig#enabledEventTypes()} are dropped.
- *
- * {@link #emitEvent(ModelUsage)} requires a non-null {@link ModelUsage} — callers must guarantee
- * usage data exists before invoking. The region for each event is read from {@link
- * ModelUsage#tenant()} ({@code Tenant.region()}).
- */
-public class DefaultBilling implements Billing {
-
- private static final Logger BILLING_LOGGER = LoggerFactory.getLogger("billing.events");
- private static final Logger LOGGER = LoggerFactory.getLogger(DefaultBilling.class);
-
- private static final ObjectWriter OBJECT_WRITER = new ObjectMapper().writer();
-
- private final String product;
- private final String resourceType;
- private final Set NOTE: we first made billing for findAndRerank / rerank models, these tests ony use
+ * $vectorize with insert, not findAndRerank, because we have test infra for a custom vectorizer
+ *
+ * TestMethodOrder is used because the second test turns off the S3 container so we need that to
+ * be last
+ */
+@QuarkusIntegrationTest
+@WithTestResource(value = DseTestResource.class)
+@WithTestResource(value = S3MockTestResource.class)
+@TestMethodOrder(MethodOrderer.OrderAnnotation.class) // needed because we turn off mock S3
+public class BillingS3UploadIntegrationTest extends AbstractCollectionIntegrationTestBase {
+
+ private static final Logger LOGGER =
+ LoggerFactory.getLogger(BillingS3UploadIntegrationTest.class);
+ private static final ObjectMapper MAPPER = new ObjectMapper();
+
+ private final TestConstants TEST_CONSTANTS = new TestConstants();
+
+ private static final int DOCUMENT_COUNT = 10;
+ private static final Pattern KEY_PATTERN =
+ Pattern.compile("data-api/\\d{4}/\\d{2}/\\d{2}/\\d{2}/\\d{2}/[0-9a-f-]{36}\\.jsonl");
+ private static final Set NOTE: must run last so because the S3 resource is shared and this tests turns if off
+ * !
+ */
+ @Test
+ @Order(Integer.MAX_VALUE)
+ public void exportFailureDoesNotAffectTheApi() {
+ S3MockTestResource.stopContainer();
+
+ // assert happens in the function, this is the test, could we insert docs ?
+ insertDocs("exportFailureDoesNotAffectTheApi() - 1st");
+
+ // wait until we register a failed upload
+ assertMetricTotal("billing.s3.failed.batches", (a) -> a.isGreaterThan(0));
+ assertMetricTotal("billing.s3.failed.events", (a) -> a.isGreaterThan(0));
+
+ // double check API still works
+ insertDocs("exportFailureDoesNotAffectTheApi() - 2nd");
+ }
+
+ // ============================================================
+ // Scaffold
+ // ============================================================
+
+ private void createCollection() {
+ var collectionOptions =
+ """
+ {
+ "vector": {
+ "metric": "cosine",
+ "dimension": 5,
+ "service": {
+ "provider": "custom",
+ "modelName": "text-embedding-ada-002",
+ "authentication": {
+ "providerKey" : "shared_creds.providerKey"
+ },
+ "parameters": {
+ "projectId": "test project"
+ }
+ }
+ }
+ }
+ """;
+
+ assertNamespaceCommand(TEST_CONSTANTS.KEYSPACE_NAME)
+ .templated()
+ .createCollection(TEST_CONSTANTS.COLLECTION_NAME, collectionOptions)
+ .wasSuccessful();
+ }
+
+ private void insertDocs(String idPrefix) {
+ for (int i = 0; i < DOCUMENT_COUNT; i++) {
+ var id = "doc-%s-%d".formatted(idPrefix, i);
+ var doc =
+ """
+ {
+ "_id": "doc-%s",
+ "description": "billing export test document %s",
+ "$vectorize": "billing export test document %s"
+ }"""
+ .formatted(id, id, id);
+ assertTableCommand(TEST_CONSTANTS.KEYSPACE_NAME, TEST_CONSTANTS.COLLECTION_NAME)
+ .templated()
+ .insertOne(doc)
+ .wasSuccessful();
+ }
+ }
+
+ private void assertIsSet(JsonNodeType nodeType, JsonNode parent, List TODO: out of order log records gets correct oldest metric TODO: TEST a big line bigger than
+ * the max bytes gets through
+ */
+public class BatchedLogBufferTest extends BillingTestBase {
+
+ // *********************************************************
+ // Offer - Producer side of the buffer
+ // *********************************************************
+
+ /** When the buffer reaches capacity calling offer() fails. Single producer thread. */
+ @Test
+ public void offerFailsAtCapacitySingleThread() {
+
+ var fixture = defaultBufferFixture(false);
+ var slice = Slice.to(BUFFER_CAPACITY);
+
+ // send full capacity to the buffer, should all work
+ fixture.assertOffer("offerFailsAtCapacitySingleThread() - prefill to capacity", slice);
+
+ // 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 = 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));
+ }
+ }
+
+ // 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 = 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);
+ 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 = defaultBufferFixture(false);
+
+ assertThatThrownBy(() -> fixture.buffer().offer(null))
+ .as("offerNullRecord() null log record is an exception")
+ .isInstanceOf(NullPointerException.class);
+ }
+
+ @Test
+ public void offerNullOrBlankMessage() {
+ var fixture = defaultBufferFixture(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 = defaultBufferFixture(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 = defaultBufferFixture(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 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 = defaultBufferFixture(false);
+
+ // Setup a Consumer thread, it will keep running until we set consumerShutdown
+ var normalBatches = new ArrayList See also {@link DefaultBillingTest}
*/
class BillingTest {
- private final TestConstants testConstants = new TestConstants();
+ private final TestConstants TEST_CONSTANTS = new TestConstants();
// ============================================================
// NO_OP behavior
@@ -37,10 +38,12 @@ void noOpDoesNothingForValidUsage() {
assertThatCode(() -> Billing.NO_OP.emitEvent(stubUsage())).doesNotThrowAnyException();
}
+ /**
+ * Even the NO-OP enforces the non-null contract so callers can't accidentally pass null and have
+ * it silently swallowed in a test that uses the no-op.
+ */
@Test
void noOpThrowsOnNullUsage() {
- // Even the NO-OP enforces the non-null contract so callers can't accidentally pass null and
- // have it silently swallowed in a test that uses the no-op.
assertThatThrownBy(() -> Billing.NO_OP.emitEvent(null))
.isInstanceOf(NullPointerException.class)
.hasMessageContaining("modelUsage");
@@ -71,39 +74,32 @@ void createReturnsNoOpWhenFeatureDisabled() {
/**
* If BILLING_EVENTS_LOGGING is enabled in startup config, a request header MUST NOT be able to
- * turn it off — the config is authoritative. Verified at the dispatch layer because that's where
- * the user-visible effect lands (you get DefaultBilling, not NO_OP).
+ * turn it off.
*/
@Test
void createConfigEnabledIsNotOverriddenByHeader() {
- var config = mock(FeaturesConfig.class);
- when(config.flags()).thenReturn(Map.of(ApiFeature.BILLING_EVENTS_LOGGING, "true"));
- var headers = MultiMap.caseInsensitiveMultiMap();
- headers.add(ApiFeature.BILLING_EVENTS_LOGGING.httpHeaderName(), "false");
- var apiFeatures =
- ApiFeatures.fromConfigAndRequest(config, new RequestContext.HttpHeaderAccess(headers));
+ // enabled in the config
+ // disabled by the header
+ var apiFeatures = featuresWithBilling(true, false);
var billing = Billing.create(validConfig(), apiFeatures);
assertThat(billing)
- .as("config=true must win over header=false at dispatch")
+ .as("config=true must win over header=false")
.isInstanceOf(DefaultBilling.class);
}
/**
* Conversely, if startup config leaves the flag unset, a request header CAN enable it. Proves the
- * header path is alive — without this, the test above would pass trivially.
+ * header path is alive
*/
@Test
void createHeaderEnablesWhenConfigUnset() {
- var config = mock(FeaturesConfig.class);
- when(config.flags()).thenReturn(Map.of());
- var headers = MultiMap.caseInsensitiveMultiMap();
- headers.add(ApiFeature.BILLING_EVENTS_LOGGING.httpHeaderName(), "true");
- var apiFeatures =
- ApiFeatures.fromConfigAndRequest(config, new RequestContext.HttpHeaderAccess(headers));
+ // no value in the config
+ // enabled by the header
+ var apiFeatures = featuresWithBilling(null, true);
var billing = Billing.create(validConfig(), apiFeatures);
@@ -116,19 +112,23 @@ void createHeaderEnablesWhenConfigUnset() {
/** Minimal valid {@link BillingConfig} — enough for {@link DefaultBilling} to construct. */
private static BillingConfig validConfig() {
+
var config = mock(BillingConfig.class);
when(config.product()).thenReturn("serverless");
when(config.resourceType()).thenReturn("serverless_database");
when(config.internalModelProviders()).thenReturn(List.of("nvidia"));
when(config.enabledEventTypes()).thenReturn(Optional.empty());
+
return config;
}
private static ApiFeatures featuresWithBilling(boolean enabled) {
- var config = mock(FeaturesConfig.class);
- when(config.flags())
- .thenReturn(Map.of(ApiFeature.BILLING_EVENTS_LOGGING, String.valueOf(enabled)));
- return ApiFeatures.fromConfigAndRequest(config, null);
+ return featuresWithBilling(enabled, null);
+ }
+
+ private static ApiFeatures featuresWithBilling(Boolean configEnabled, Boolean headerEnabled) {
+ return ApiFeaturesTestUtil.withFeature(
+ ApiFeature.BILLING_EVENTS_LOGGING, configEnabled, headerEnabled);
}
private ModelUsage stubUsage() {
@@ -136,7 +136,7 @@ private ModelUsage stubUsage() {
ModelProvider.NVIDIA,
ModelType.EMBEDDING,
"model-x",
- testConstants.TENANT,
+ TEST_CONSTANTS.TENANT,
ModelInputType.INDEX,
100,
500,
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..38bceb85f6
--- /dev/null
+++ b/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingTestBase.java
@@ -0,0 +1,520 @@
+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.MetricsUnitAssertions;
+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 buffered and uploaded */
+public abstract class BillingTestBase {
+
+ private static final Logger LOGGER = LoggerFactory.getLogger(BillingTestBase.class);
+
+ // ======================================================================
+ // BUFFER config and values for the tests to use
+ // NOTE: these values are all balanced, see comments before changing
+ // ======================================================================
+
+ // 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 ...
+ */
+ record Slice(int from, int to) {
+
+ public ...
+ */
+ record BufferSnapshot(
+ boolean isEmpty,
+ int size,
+ long queuedBytes,
+ int remainingCapacity,
+ Fixture fixture,
+ MetricsUnitAssertions.MetricSnapshot metricSnapshot) {
+
+ static BufferSnapshot create(Fixture fixture) {
+
+ return new BufferSnapshot(
+ fixture.buffer().isEmpty(),
+ fixture.buffer().size(),
+ fixture.buffer().queuedBytes(),
+ fixture.buffer().remainingCapacity(),
+ fixture,
+ fixture.metricsUtil.createSnapshot());
+ }
+
+ /**
+ * 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());
+ // metric should be the size of the buffer, does not matter what metric snapshot was
+ fixture.metricsUtil.assertMetric(fixture.bufferMetrics.size, size() + slice.size());
+
+ // Buffer offered increases, dropped does not change
+ fixture.metricsUtil.assertMetric(metricSnapshot, fixture.bufferMetrics.offered, slice.size());
+ fixture.metricsUtil.assertMetric(metricSnapshot, fixture.bufferMetrics.dropped, 0);
+
+ // remaining is just a calculation, capacity - size
+ fixture.metricsUtil.assertMetric(
+ fixture.bufferMetrics.remainingCapacity, fixture.queueCapacity - fixture.buffer().size());
+
+ // head age is the age of the first item in the buffer, NOTE clock overrides can change time
+ var firstRecord = fixture.buffer().peekBuffer().getFirst();
+ if (firstRecord != null) {
+ var clockForMetric =
+ fixture.clock() == null ? BatchedLogBuffer.DEFAULT_CLOCK : fixture.clock();
+ // metric is calc's when read, so it may be younger (less ms) but never older (more ms)
+ fixture.metricsUtil.assertMetric(
+ fixture.bufferMetrics.headAgeMs,
+ (a) -> {
+ a.isLessThanOrEqualTo(
+ Duration.between(firstRecord.eventAt(), clockForMetric.instant())
+ .abs()
+ .toMillis());
+ });
+ }
+
+ 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);
+ fixture.metricsUtil.assertMetric(fixture.bufferMetrics.bytes, 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..d8931b5166
--- /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;
+
+/** Tests for the {@link BillingUploadingLogHandler} */
+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 These tests only check it calls the S3 client as expected, not that batches are sent anywhere.
+ * See {@link BillingS3UploadIntegrationTest}
+ */
+class S3BatchedLogUploaderTest {
+
+ private static final Instant OLDEST_EVENT_AT = Instant.parse("2026-05-20T16:23:00Z");
+ private static final String LINE_A = "{\"a\":1}"; // him
+ private static final String LINE_B = "{\"b\":2}";
+ private static final int LINE_BYTE_SIZE = 16; // 7 chars per line and one for new line
+ private static final BatchedLogBuffer.Batch BATCH =
+ new BatchedLogBuffer.Batch(
+ BatchedLogBuffer.BillingBatchReason.DRAINING,
+ List.of(LINE_A, LINE_B),
+ LINE_BYTE_SIZE, // 7 chars per line and one for new line
+ OLDEST_EVENT_AT);
+ private static final String EXPECTED_KEY =
+ "data-api/2026/05/20/16/23/%s.jsonl".formatted(BATCH.id());
+
+ @Test
+ void createReadsConfig() {
+
+ var fixture = Fixture.create(true);
+
+ var toString = fixture.uploader.toString();
+
+ assertThat(toString)
+ .as("toString() contains region)")
+ .contains("region=" + fixture.config.region());
+ assertThat(toString)
+ .as("toString() contains bucket)")
+ .contains("bucket=" + fixture.config.bucket());
+ assertThat(toString)
+ .as("toString() contains pathPrefix)")
+ .contains("pathPrefix=" + fixture.config.s3PathPrefix());
+
+ // TODO: test for endpoint override.
+ }
+
+ @Test
+ void closeCascadedToS3Client() {
+
+ var fixture = Fixture.create();
+
+ fixture.uploader.close();
+ verify(fixture.s3Client, times(1)).close();
+ }
+
+ /** ObjectKey is buikld using the properties of the Batch */
+ @Test
+ void objectLocationUsesBatchProperties() {
+
+ var fixture = Fixture.create();
+
+ var location = fixture.uploader.objectLocation(BATCH);
+ assertThat(location.region()).as("region is same as config").isEqualTo(fixture.config.region());
+
+ assertThat(location.bucket()).as("bucket is same as config").isEqualTo(fixture.config.bucket());
+
+ assertThat(location.key()).as("key is as expected").isEqualTo(EXPECTED_KEY);
+ }
+
+ /** Object body built using properties of the Batch */
+ @Test
+ void objectContentContainsBatchLines() {
+
+ var fixture = Fixture.create();
+
+ var bytes = fixture.uploader.objectContent(BATCH);
+
+ assertThat(bytes)
+ .as("bytes is as expected")
+ .isEqualTo((LINE_A + "\n" + LINE_B + "\n").getBytes(StandardCharsets.UTF_8));
+ }
+
+ /** Request to S3 clent matches the batch and config properties */
+ @Test
+ void uploadUsesCorrectBatchProperties() {
+ var fixture = Fixture.create();
+
+ fixture.uploader.upload(BATCH);
+
+ var requestCaptor = ArgumentCaptor.forClass(PutObjectRequest.class);
+ var bodyCaptor = ArgumentCaptor.forClass(AsyncRequestBody.class);
+ verify(fixture.s3Client, times(1)).putObject(requestCaptor.capture(), bodyCaptor.capture());
+
+ var request = requestCaptor.getValue();
+ assertThat(request.bucket()).as("bucket is same as config").isEqualTo(fixture.config.bucket());
+ assertThat(request.key()).as("key is as expected").isEqualTo(EXPECTED_KEY);
+ assertThat(request.contentType())
+ .as("content type is as expected")
+ .isEqualTo(S3BatchedLogUploader.CONTENT_TYPE_NDJSON);
+
+ var body = bodyCaptor.getValue();
+ assertThat(body.contentLength()).as("content length matches batch").hasValue(BATCH.bytes());
+ }
+
+ /** Multiple calls to upload generate different keys, each matching batch properties */
+ @Test
+ void uploadGeneratesUniqueKeys() {
+ var fixture = Fixture.create();
+
+ // second batch we will send, has diff Id and 1 hour later.
+ var oldestEvent2 = OLDEST_EVENT_AT.plus(1, ChronoUnit.HOURS);
+ var batch2 =
+ new BatchedLogBuffer.Batch(
+ BatchedLogBuffer.BillingBatchReason.MAX_BYTES_EXCEEDED,
+ BATCH.lines(),
+ LINE_BYTE_SIZE,
+ oldestEvent2);
+
+ fixture.uploader.upload(BATCH);
+ fixture.uploader.upload(batch2);
+
+ var requestCaptor = ArgumentCaptor.forClass(PutObjectRequest.class);
+ var bodyCaptor = ArgumentCaptor.forClass(AsyncRequestBody.class);
+ verify(fixture.s3Client, times(2)).putObject(requestCaptor.capture(), bodyCaptor.capture());
+
+ var requests = requestCaptor.getAllValues();
+
+ assertThat(requests.stream().map(PutObjectRequest::key))
+ .as("all put requests keys are unique")
+ .doesNotHaveDuplicates();
+
+ assertThat(requests.getFirst().key()).as("key for batch 1 as expected").isEqualTo(EXPECTED_KEY);
+
+ var expectedKey2 = "data-api/2026/05/20/17/23/%s.jsonl".formatted(batch2.id());
+ assertThat(requests.get(1).key()).as("key for batch 2 as expected").isEqualTo(expectedKey2);
+ }
+
+ @Test
+ void uploadS3FailureReturned() {
+
+ var fixture = Fixture.create();
+ var expectedThrowable = new RuntimeException("S3 Failed");
+ when(fixture.s3Client.putObject(any(PutObjectRequest.class), any(AsyncRequestBody.class)))
+ .thenReturn(CompletableFuture.failedFuture(expectedThrowable));
+
+ var uploadResult = fixture.uploader.upload(BATCH).await().atMost(Duration.ofSeconds(10));
+
+ verify(
+ fixture.s3Client,
+ times(1)
+ .description("s3Client upload() called once, all retry is internal to s3 client"))
+ .putObject(any(PutObjectRequest.class), any(AsyncRequestBody.class));
+
+ assertThat(uploadResult).isNotNull();
+ assertThat(uploadResult.throwable())
+ .as("uploadResult has expected throwable instance")
+ .isSameAs(expectedThrowable);
+ assertThat(uploadResult.batch()).as("uploadResult has expected batch instance").isSameAs(BATCH);
+ }
+
+ @Test
+ void uploadS3SuccessReturned() {
+
+ var fixture = Fixture.create();
+
+ var uploadResult = fixture.uploader.upload(BATCH).await().atMost(Duration.ofSeconds(10));
+
+ verify(
+ fixture.s3Client,
+ times(1)
+ .description("s3Client upload() called once, all retry is internal to s3 client"))
+ .putObject(any(PutObjectRequest.class), any(AsyncRequestBody.class));
+
+ assertThat(uploadResult).isNotNull();
+ assertThat(uploadResult.throwable()).as("uploadResult throwable is null when success").isNull();
+ assertThat(uploadResult.batch()).as("uploadResult has expected batch instance").isSameAs(BATCH);
+ }
+
+ // ============================================================
+ // Scaffold
+ // ============================================================
+
+ private record Fixture(
+ S3BatchedLogUploader uploader,
+ S3AsyncClient s3Client,
+ PutObjectResponse putResponse,
+ BatchedLogUploaderMetrics metrics,
+ BillingS3UploadConfig config) {
+
+ static Fixture create() {
+ return create(false);
+ }
+
+ static Fixture create(boolean useFactory) {
+
+ var config =
+ new SmallRyeConfigBuilder()
+ .withMapping(BillingS3UploadConfig.class)
+ .build()
+ .getConfigMapping(BillingS3UploadConfig.class);
+
+ var metrics = mock(BatchedLogUploaderMetrics.class);
+
+ var s3Client = mock(S3AsyncClient.class);
+
+ var sdkResponse = SdkHttpResponse.builder().statusCode(200).build();
+ var builder = PutObjectResponse.builder();
+ // sdkHttpResponse is on an inherited builder, return type changes
+ builder.sdkHttpResponse(sdkResponse);
+ var putResponse = builder.build();
+
+ when(s3Client.putObject(any(PutObjectRequest.class), any(AsyncRequestBody.class)))
+ .thenReturn(CompletableFuture.completedFuture(putResponse));
+
+ S3BatchedLogUploader uploader;
+ if (useFactory) {
+ uploader =
+ S3BatchedLogUploader.create(
+ config.region(),
+ config.bucket(),
+ config.endpointOverride().orElse(null),
+ config.s3PathPrefix(),
+ config.s3CallAttemptTimeout(),
+ config.s3TotalCallTimeout(),
+ config.s3RetryMode(),
+ metrics);
+ } else {
+ uploader =
+ new S3BatchedLogUploader(
+ s3Client, config.region(), config.bucket(), config.s3PathPrefix(), metrics);
+ }
+
+ return new Fixture(uploader, s3Client, putResponse, metrics, config);
+ }
+ }
+}
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..8ac254e887
--- /dev/null
+++ b/src/test/java/io/stargate/sgv2/jsonapi/testresource/S3MockTestResource.java
@@ -0,0 +1,112 @@
+package io.stargate.sgv2.jsonapi.testresource;
+
+import static io.stargate.sgv2.jsonapi.util.SmallRyeConfigTestUtil.addPropertyTo;
+
+import com.adobe.testing.s3mock.testcontainers.S3MockContainer;
+import io.quarkus.test.common.QuarkusTestResourceLifecycleManager;
+import io.stargate.sgv2.jsonapi.config.BillingS3UploadConfig;
+import java.util.HashMap;
+import java.util.Map;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Starts an S3Mock container (https://github.com/adobe/S3Mock) and configures billing to send
+ * events with short pauses.
+ */
+public class S3MockTestResource implements QuarkusTestResourceLifecycleManager {
+
+ private static final Logger LOGGER = 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 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 static volatile 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;
+ }
+
+ /**
+ * Stops the S3Mock container, leaving nothing listening on the exported endpoint: every upload
+ * from then on fails with connection-refused.
+ */
+ public static void stopContainer() {
+ if (container == null) {
+ throw new IllegalStateException("S3MockTestResource has not been started");
+ }
+ container.stop();
+ }
+
+ @Override
+ public Map Example usage:
+ *
+ * Example:
+ *
+ * Create a single instance for each test, and then use the assertion methods to check metric
+ * values.
+ *
+ * NOTE: because metrics like a {@link Counter} can only be increased, and the only way to
+ * reset the metric values is to destroy and re-creat. Use the {@link MetricSnapshot} created with
+ * {@link #createSnapshot()} to assert that a metric has changed a certain amount *since* the
+ * snapshot was taken.
+ */
+public class MetricsUnitAssertions {
+
+ private final MeterRegistry meterRegistry;
+
+ public MetricsUnitAssertions() {
+ this(new SimpleMeterRegistry());
+ }
+
+ public MetricsUnitAssertions(MeterRegistry meterRegistry) {
+ this.meterRegistry = Objects.requireNonNull(meterRegistry, "meterRegistry is null");
+ }
+
+ public MeterRegistry registry() {
+ return meterRegistry;
+ }
+
+ public MetricSnapshot createSnapshot() {
+
+ Map See {@link #propertyName(ConfigMethodRef)} and {@link #addPropertyTo(Map, ConfigMethodRef,
+ * String)}
+ */
+public final class SmallRyeConfigTestUtil {
+
+ /**
+ * Functional interface to represent a method on a SmallRye config interface.
+ *
+ * {@link Serializable} is required, the java compiler only emits the synthetic Used with {@link #propertyName(ConfigMethodRef)} , the T type param is the instance that
+ * would be passed if the Function was evaluated. This is why it's a Function not a Supplier even
+ * though the Config function takes no params.
+ *
+ * You will probably want to use {@link #addPropertyTo(Map, ConfigMethodRef, String)}
+ *
+ * Example, given this config:
+ *
+ * Example usage:
+ *
+ * true billing events are sent to S3 using the config below. */
+ @WithDefault("false")
+ boolean enabled();
+
+ /**
+ * When true adding the S3 handler for logging at run time removes any other handler
+ * the logger has installed.
+ *
+ *
+ * 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 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 BlockingQueuenull if there is no next batch. Keep calling with drainFully true
+ * until a null is returned to drain the buffer.
+ */
+ public Batch nextBatch(boolean drainFully) {
+
+ var batchReason = decideNextBatch(drainFully);
+ if (batchReason == null) {
+ return null;
+ }
+
+ List
- *
- *
- * Pass each aggregated {@link ModelUsage} to {@link #emitEvent(ModelUsage)}. {@code modelUsage}
- * must be non-null — callers are responsible for ensuring usage data exists before invoking.
+ *
+ *
*
* @param config billing configuration; only consulted when the feature is enabled
* @param apiFeatures the request's resolved feature set; must not be null
*/
static Billing create(BillingConfig config, ApiFeatures apiFeatures) {
+
Objects.requireNonNull(apiFeatures, "apiFeatures must not be null");
+
return apiFeatures.isFeatureEnabled(ApiFeature.BILLING_EVENTS_LOGGING)
? new DefaultBilling(config)
: NO_OP;
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..9ec7e70e45 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;
@@ -13,7 +13,7 @@
* A billing event emitted as a structured JSON log line for downstream billing pipelines.
*
*
+ *
+ */
+@ApplicationScoped
+public class BillingS3Lifecycle {
+
+ private static final org.slf4j.Logger LOGGER = LoggerFactory.getLogger(BillingS3Lifecycle.class);
+
+ public static final String METRICS_PREFIX = "billing";
+
+ private final BillingS3UploadConfig config;
+ private final MeterRegistry meterRegistry;
+
+ private volatile BillingUploadingLogHandler handler;
+
+ @Inject
+ public BillingS3Lifecycle(BillingS3UploadConfig config, MeterRegistry meterRegistry) {
+ this.config = config;
+ this.meterRegistry = meterRegistry;
+ }
+
+ @VisibleForTesting
+ protected BillingUploadingLogHandler handler() {
+ return this.handler;
+ }
+
+ void onStart(@Observes StartupEvent event) {
+
+ if (!config.enabled()) {
+ LOGGER.info("onStart() - S3 export disabled");
+ return;
+ }
+
+ LOGGER.info("onStart() - S3 export enabled");
+
+ var uploader =
+ S3BatchedLogUploader.create(
+ config.region(),
+ config.bucket(),
+ config.endpointOverride().orElse(null),
+ config.s3PathPrefix(),
+ config.s3CallAttemptTimeout(),
+ config.s3TotalCallTimeout(),
+ config.s3RetryMode(),
+ new BatchedLogUploaderMetrics(meterRegistry, METRICS_PREFIX + ".s3"));
+ LOGGER.info("onStart() - using uploader: {}", uploader);
+
+ var buffer =
+ new BatchedLogBuffer(
+ config.bufferMaxBatchSize(),
+ config.bufferMaxBatchBytes(),
+ config.bufferMaxBatchAge(),
+ config.bufferCapacity(),
+ new BatchedLogBufferMetrics(meterRegistry, METRICS_PREFIX));
+ LOGGER.info("onStart() - using log buffer: {}", buffer);
+
+ this.handler =
+ new BillingUploadingLogHandler(
+ buffer,
+ uploader,
+ config.handlerSleepDuration(),
+ config.handlerUploadSafetyDeadline(),
+ config.handlerUploadShutdownDeadline());
+ LOGGER.info("onStart() - using handler: {}", handler);
+
+ var billingLogger = Logger.getLogger(DefaultBilling.BILLING_LOGGER_NAME);
+ if (config.disableOtherHandlers()) {
+ LOGGER.info("onStart() - removing existing log handlers");
+ for (var existing : billingLogger.getHandlers()) {
+ LOGGER.info("onStart() - removing existing log handler. existing:{} ", existing);
+ billingLogger.removeHandler(existing);
+ }
+ } else {
+ LOGGER.info("onStart() - leaving existing log handlers");
+ }
+
+ billingLogger.addHandler(this.handler);
+ LOGGER.info(
+ "onStart() - attached log handler to logger. BILLING_LOGGER_NAME: {}",
+ DefaultBilling.BILLING_LOGGER_NAME);
+
+ Infrastructure.getDefaultWorkerPool()
+ .execute(
+ () -> {
+ LOGGER.info(
+ "onStart() - on worked pool thread, calling handler.startUploading() on this thread");
+ this.handler.startUploading();
+ });
+ }
+
+ void onStop(@Observes ShutdownEvent event) {
+
+ if (this.handler == null) {
+ LOGGER.info("onStop() - handler was null, nothing to close");
+ return;
+ }
+
+ Logger.getLogger(DefaultBilling.BILLING_LOGGER_NAME).removeHandler(this.handler);
+ LOGGER.info(
+ "onStop() - handler removed from logger. BILLING_LOGGER_NAME:{}",
+ DefaultBilling.BILLING_LOGGER_NAME);
+
+ // 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 {
+ LOGGER.debug("onStop() - calling handler.close()");
+ this.handler.close();
+ LOGGER.info("onStop() - handler was closed without error");
+ } catch (Exception e) {
+ LOGGER.warn("onStop() - error calling handler.close(), swallowing", e);
+ }
+ }
+}
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..7d3091ba25
--- /dev/null
+++ b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingUploadingLogHandler.java
@@ -0,0 +1,382 @@
+package io.stargate.sgv2.jsonapi.service.billing;
+
+import static io.stargate.sgv2.jsonapi.util.ClassUtils.classSimpleName;
+
+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.
+ *
+ *
+ *
+ */
+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 running. 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;
+
+ /**
+ * Creates a new instance, see {@link BillingS3Lifecycle}
+ *
+ * @param buffer A configured buffer that is used to buffer incoming log messages
+ * @param uploader {@link AsyncBatchedLogUploader} to upload messages when the buffer has a new
+ * batch.
+ * @param uploadSleepDuration Duration the thread running {@link #startUploading()} will sleep
+ * waiting for a new batch from the buffer.
+ * @param uploaderSafetyDeadline Duration to wait for the uploader to return from uploading. NOTE:
+ * the uploader should have its own internal timeouts, this is to prevent the uploading thread
+ * blocking.
+ * @param uploadShutdownDeadline Duration a call to {@link #close()} should wait for the uploading
+ * thread to complete flushing remaining events.
+ */
+ 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 publishes 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.
+ *
+ * {@code
+ * ApiFeatures features = ApiFeaturesTestUtil.withFeature(
+ * ApiFeature.BILLING_EVENTS_LOGGING, configEnabled, headerEnabled);
+ *
+ * }
+ *
+ * @param feature The {@link ApiFeature} we want to set value for
+ * @param configEnabled If the feature should be set via config, if null is ignored.
+ * @param headerEnabled If the feature should be set via HTTP Header, if null is ignored.
+ * @return Configured {@link ApiFeatures}
+ */
+ public static ApiFeatures withFeature(
+ ApiFeature feature, Boolean configEnabled, Boolean headerEnabled) {
+
+ Map{@code
+ * assertMetricTotal("billing.s3.uploaded.events", 10);
+ * }
+ */
+ public static void assertMetricTotal(String metricName, double metricValue) {
+ assertMetricTotal(metricName, metricValue, DEFAULT_AWAIT_DURATION, DEFAULT_POLL_DURATION);
+ }
+
+ /**
+ * Example:
+ *
+ * {@code
+ * assertMetricTotal("billing.s3.uploaded.events", (a) -> a.isGreaterThanOrEqualTo(10));
+ * }
+ */
+ public static void assertMetricTotal(
+ String metricName, ConsumerassertConsumer */
+ public static void assertMetricTotal(
+ String metricName,
+ Consumer
+ * writeReplace() function for a lambda whose target type is serializable.
+ *
+ *
+ * SmallRyeConfigUtil.propertyName(MyConfigInterface::host);
+ *
+ *
+ * @param this if the function was invoked.
+ * @param {@code
+ * @ConfigMapping(prefix = "stargate.jsonapi.billing")
+ * public interface BillingConfig {
+ * @WithDefault("serverless")
+ * String product();
+ * }
+ *
+ * }
+ *
+ * Invoked using:
+ *
+ * {@code
+ * SmallRyeConfigUtil.propertyName(BillingConfig::product);
+ *
+ * }
+ *
+ * Will return the string:
+ *
+ * {@code
+ * stargate.jsonapi.billing.product
+ *
+ * }
+ */
+ public static {@code
+ * Map
+ */
+ public static