queue;
+ private final AtomicLong queuedBytes = new AtomicLong(0);
+ private final BatchedLogBufferMetrics metrics;
+
+ /** 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 many 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 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 incase 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.offered();
+ if (!queue.offer(newEntry)) {
+ // Bounded buffer full, drop and count
+ LOGGER.debug("offer() - buffer full, dropping new entry: {}", newEntry);
+ metrics.dropped();
+ 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
+ * null if there is no next batch.
+ */
+ public Batch nextBatch(boolean drainFully) {
+
+ var batchReason = decideNextBatch(drainFully);
+ if (batchReason == null) {
+ return null;
+ }
+
+ List batchLines = new ArrayList<>(maxBatchSize);
+ Instant oldestEventAt = null;
+ long batchBytes = 0;
+ Entry peeked;
+
+ // No matter why we started we create a full batch, e.g. we could start because the oldest
+ // entry is past maxAge, but we still fill the batch.
+ while (batchLines.size() < maxBatchSize && ((peeked = queue.peek()) != null)) {
+
+ var lineBytes = peeked.lineBytes();
+ if (batchBytes + lineBytes > maxBatchBytes && !batchLines.isEmpty()) {
+ // adding the next line will be too many bytes, we can only do this if the batch
+ // is empty, so a single big message can be put into a batch and not block everyone else
+ // break out of here.
+ break;
+ }
+
+ // OK to remove entry from buffer and add to batch
+ // there is only this thread as a consumer, no race condition
+ var polled = queue.poll();
+ // sanity check
+ if (polled != peeked) {
+ throw new IllegalStateException(
+ "nextBatch() - peeked entry is not same object as polled entry");
+ }
+ if (oldestEventAt == null || polled.eventAt().isBefore(oldestEventAt)) {
+ oldestEventAt = polled.eventAt();
+ }
+ batchLines.add(polled.line());
+ queuedBytes.addAndGet(-lineBytes);
+ batchBytes += lineBytes;
+ }
+
+ if (batchLines.isEmpty() && !queue.isEmpty()) {
+ // sanity check
+ // there is messages in the queue, but we did not put any in the batch, something wrong
+ // but it could be a race condition - things may be added after the loop finish
+ // so do not throw, just log
+ LOGGER.warn(
+ "nextBatch() - did not add any lines for next batch, but queue is not empty. May be logic bug or expected race condition. queue.size:{}",
+ queue.size());
+ return null;
+ }
+
+ if (LOGGER.isDebugEnabled()) {
+ LOGGER.debug(
+ "nextBatch() - next batch created, reason:{}, batchLines.size:{}, batchBytes:{}, oldestEventAt: {}",
+ batchReason,
+ batchLines.size(),
+ batchBytes,
+ oldestEventAt);
+ }
+ return new Batch(batchReason, batchLines, batchBytes, oldestEventAt, clock);
+ }
+
+ /** Gets a copy of the contents of the buffer in a new array list, for testing. */
+ @VisibleForTesting
+ List peekBuffer() {
+ return new ArrayList<>(queue);
+ }
+
+ public boolean isEmpty() {
+ return queue.isEmpty();
+ }
+
+ public int size() {
+ return queue.size();
+ }
+
+ public long queuedBytes() {
+ return queuedBytes.get();
+ }
+
+ public int remainingCapacity() {
+ return queue.remainingCapacity();
+ }
+
+ /**
+ * Gets the age of the item at the head of the buffer.
+ *
+ * Age is determined by the clock used to create the buffer.
+ *
+ * @return age of the head item in the buffer, or null if no items in the buffer.
+ */
+ public Duration headEntryAge() {
+ return entryAge(queue.peek());
+ }
+
+ @VisibleForTesting
+ Duration entryAge(Entry entry) {
+ return entry == null ? Duration.ZERO : Duration.between(entry.eventAt(), clock.instant());
+ }
+
+ private BillingBatchReason decideNextBatch(boolean drainFully) {
+
+ 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 lines;
+ private final long bytes;
+ private final Instant oldestEventAt;
+ private final Clock clock;
+
+ Batch(BillingBatchReason reason, List lines, long bytes, Instant oldestEventAt) {
+ this(reason, lines, bytes, oldestEventAt, DEFAULT_CLOCK);
+ }
+
+ private Batch(
+ BillingBatchReason reason,
+ List lines,
+ long bytes,
+ Instant oldestEventAt,
+ Clock clock) {
+ this.reason = Objects.requireNonNull(reason, "reason must not be null");
+ this.lines =
+ Collections.unmodifiableList(Objects.requireNonNull(lines, "lines must not be null"));
+ this.bytes = bytes;
+ this.oldestEventAt = Objects.requireNonNull(oldestEventAt, "oldestEventAt must not be null");
+ this.clock = Objects.requireNonNull(clock, "clock must not be null");
+ }
+
+ public UUID id() {
+ return id;
+ }
+
+ public BillingBatchReason reason() {
+ return reason;
+ }
+
+ /**
+ * @return Unmodifiable list of the log lines in this buffer
+ */
+ public List lines() {
+ return lines;
+ }
+
+ public Instant oldestEventAt() {
+ return oldestEventAt;
+ }
+
+ public Duration oldestEventAtDuration() {
+ // using the outer buffers clock, so any tests that change the clock
+ // get a consistent results
+ return Duration.between(oldestEventAt(), clock.instant());
+ }
+
+ public int size() {
+ return lines.size();
+ }
+
+ public long bytes() {
+ return bytes;
+ }
+
+ @Override
+ public String toString() {
+ return new StringBuilder(classSimpleName(this) + "{")
+ .append("id=")
+ .append(id)
+ .append(", reason=")
+ .append(reason)
+ .append(", oldestEventAt=")
+ .append(oldestEventAt)
+ .append(", size=")
+ .append(size())
+ .append(", bytes=")
+ .append(bytes())
+ .append("}")
+ .toString();
+ }
+ }
+
+ /**
+ * An entry in a batch, this is one log message passed to the buffer.
+ *
+ * @param eventAt When the log event happened
+ * @param line The log message
+ */
+ public record Entry(Instant eventAt, String line) {
+
+ /**
+ * Gets the length of the line in bytes,
+ *
+ * Kind of a hack, we are counting Unicode code points and calling that 1 byte. Should work
+ * for ascii text, will undercount if there is non ASCII chars but 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 95%
rename from src/main/java/io/stargate/sgv2/jsonapi/service/provider/Billing.java
rename to src/main/java/io/stargate/sgv2/jsonapi/service/billing/Billing.java
index 3bb0a72549..a0c3fd50a8 100644
--- a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/Billing.java
+++ b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/Billing.java
@@ -1,8 +1,9 @@
-package io.stargate.sgv2.jsonapi.service.provider;
+package io.stargate.sgv2.jsonapi.service.billing;
import io.stargate.sgv2.jsonapi.config.BillingConfig;
import io.stargate.sgv2.jsonapi.config.feature.ApiFeature;
import io.stargate.sgv2.jsonapi.config.feature.ApiFeatures;
+import io.stargate.sgv2.jsonapi.service.provider.ModelUsage;
import java.util.Objects;
/**
diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingEvent.java b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingEvent.java
similarity index 98%
rename from src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingEvent.java
rename to src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingEvent.java
index 2f2bbcdfa3..19dbe59263 100644
--- a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingEvent.java
+++ b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingEvent.java
@@ -1,4 +1,4 @@
-package io.stargate.sgv2.jsonapi.service.provider;
+package io.stargate.sgv2.jsonapi.service.billing;
import com.fasterxml.jackson.annotation.JsonInclude;
import com.fasterxml.jackson.annotation.JsonProperty;
diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingEventType.java b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingEventType.java
similarity index 98%
rename from src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingEventType.java
rename to src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingEventType.java
index 43f90cb2d0..c888bc1d36 100644
--- a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/BillingEventType.java
+++ b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingEventType.java
@@ -1,4 +1,4 @@
-package io.stargate.sgv2.jsonapi.service.provider;
+package io.stargate.sgv2.jsonapi.service.billing;
import com.fasterxml.jackson.annotation.JsonValue;
import java.util.EnumSet;
diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3HandlerInstaller.java b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3HandlerInstaller.java
new file mode 100644
index 0000000000..688aadbee8
--- /dev/null
+++ b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3HandlerInstaller.java
@@ -0,0 +1,139 @@
+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.BillingS3ExportConfig;
+import io.stargate.sgv2.jsonapi.metrics.BatchedLogBufferMetrics;
+import io.stargate.sgv2.jsonapi.metrics.BatchedLogUploaderMetrics;
+import jakarta.enterprise.context.ApplicationScoped;
+import jakarta.enterprise.event.Observes;
+import jakarta.inject.Inject;
+import java.util.logging.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Encapsulates setting up billing to send billing events to S3.
+ *
+ *
+ *
+ *
+ * - Encapsulates all the quarkus / CDI injection in here, so the billing classes are not bound
+ * to that approach.
+ *
- Subscribes to the quarkus lifecycle for startup and shutdown to configure billing upload,
+ * get it's uploading thread running, and then close it so we flush when shutting down.
+ *
+ */
+@ApplicationScoped
+public class BillingS3HandlerInstaller {
+
+ private static final org.slf4j.Logger LOGGER =
+ LoggerFactory.getLogger(BillingS3HandlerInstaller.class);
+
+ private static final String METRICS_PREFIX = "billing";
+ // TODO: MOVE , this is duplicated
+ public static final String BILLING_LOGGER_NAME = "billing.events";
+
+ private final BillingS3ExportConfig config;
+ private final MeterRegistry meterRegistry;
+
+ private volatile BillingUploadingLogHandler handler;
+
+ @VisibleForTesting
+ BillingUploadingLogHandler handler() {
+ return this.handler;
+ }
+
+ @Inject
+ public BillingS3HandlerInstaller(BillingS3ExportConfig config, MeterRegistry meterRegistry) {
+ this.config = config;
+ this.meterRegistry = meterRegistry;
+ }
+
+ 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));
+ LOGGER.info("onStart() - using uploader: {}", uploader);
+
+ var buffer =
+ new BatchedLogBuffer(
+ config.bufferMaxBatchSize(),
+ config.bufferMaxBatchBytes(),
+ config.bufferMaxBatchAge(),
+ config.queueCapacity(),
+ 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(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: {}", 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(BILLING_LOGGER_NAME).removeHandler(this.handler);
+ LOGGER.info(
+ "onStop() - handler removed from logger. BILLING_LOGGER_NAME:{}", 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..e11c77e6d2
--- /dev/null
+++ b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingUploadingLogHandler.java
@@ -0,0 +1,369 @@
+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.
+ *
+ * See {@link BillingS3HandlerInstaller} for setup.
+ *
+ *
// AI SLOP BELOW JUL handler that turns {@code billing.events} log lines into batched S3
+ * objects.
+ *
+ *
Division of labor: {@link BatchedLogBuffer} decides when a batch seals, {@link
+ * AsyncBatchedLogUploader} decides what an S3 object looks like, and this class decides when
+ * uploads run — the flush triggers (seal on publish, age tick, drain on close), the
+ * upload-concurrency gate, and metrics.
+ *
+ *
Delivery is at-most-once by design: publish never waits for queue capacity, full buffers drop
+ * new lines, and close drains best-effort within {@code shutdownTimeout}.
+ */
+public final class BillingUploadingLogHandler extends Handler {
+
+ // Logger for this handler, not the destination we are sending events to.
+ private static final Logger LOGGER = LoggerFactory.getLogger(BillingUploadingLogHandler.class);
+
+ /**
+ * When true means the Handler has been closed via {@link #close()} and it will silently drop any
+ * further calls to publish log entries. This also cause the upload thread to empty the buffer
+ */
+ private final AtomicBoolean isClosed = new AtomicBoolean(false);
+
+ /**
+ * Disposable permitting system for forcing wakeup in the uploading thread. startUploading() will
+ * tryAcquire() but because the permit count is 0 will always timeout, this is the timeout to wake
+ * and check buffer. When we want to force wakeup, e.g. flush(), we call release() that means any
+ * tryAcquire() returns and decrements count to 0. Resetting back to initial state. Because the
+ * wakeup permit lasts until tryAcquire it removes race conditions that could happen when
+ * flush()/notify() on an object lands before the upload thread is in wait() - if we used
+ * Object.notify() and .wait()
+ */
+ private final Semaphore wakeupPermit = new Semaphore(0);
+
+ /**
+ * There is only 1 permit for the upload process to be runnning. When {@link #startUploading()}
+ * starts it takes the permit, gives it back when the function exits (after {@link #close()}.
+ * close() uses this to make sure uploading has finished.
+ */
+ private final Semaphore uploadPermit = new Semaphore(1);
+
+ private final AsyncBatchedLogUploader uploader;
+ private final BatchedLogBuffer buffer;
+ private final Duration uploadSleepDuration;
+ private final Duration uploaderSafetyDeadline;
+ private final Duration uploadShutdownDeadline;
+
+ /** See {@link BillingS3HandlerInstaller} */
+ BillingUploadingLogHandler(
+ BatchedLogBuffer buffer,
+ AsyncBatchedLogUploader uploader,
+ Duration uploadSleepDuration,
+ Duration uploaderSafetyDeadline,
+ Duration uploadShutdownDeadline) {
+
+ this.buffer = Objects.requireNonNull(buffer, "buffer must not be null");
+ this.uploader = Objects.requireNonNull(uploader, "uploader must not be null");
+ this.uploadSleepDuration =
+ Objects.requireNonNull(uploadSleepDuration, "uploadSleepDuration must not be null");
+ this.uploaderSafetyDeadline =
+ Objects.requireNonNull(uploaderSafetyDeadline, "uploaderSafetyDeadline must not be null");
+ this.uploadShutdownDeadline =
+ Objects.requireNonNull(uploadShutdownDeadline, "uploadShutdownDeadline must not be null");
+ }
+
+ /**
+ * WARNING - sets the flag for closing but does not run the full close. just here for testing how
+ * uploading wakes up when flush called.
+ */
+ @VisibleForTesting
+ void unsafeClose() {
+ LOGGER.warn("WARNING - unsafeClose() called, must only be used in testing");
+ isClosed.set(true);
+ }
+
+ /**
+ * WARNING - acquires the upload permit, this stops the startUpload() function and close() from
+ * working normally. For testing only.
+ */
+ @VisibleForTesting
+ void unsafeAcquireUploadPermit() {
+ LOGGER.warn("WARNING - unsafeAcquireUploadPermit() called, must only be used in testing");
+ uploadPermit.acquireUninterruptibly();
+ }
+
+ // ============================================================
+ // Overrides for java.util.logging.Handler
+ // ============================================================
+
+ /**
+ * Buffers and then published the record to S3.
+ *
+ * @param record description of the log event. A null record is silently ignored and is not
+ * published
+ */
+ @Override
+ public void publish(LogRecord record) {
+
+ // Sanity check
+ if (record == null) {
+ return;
+ }
+
+ if (isClosed.get()) {
+ LOGGER.warn("publish() - called when closed, dropping record:{}", record);
+ return;
+ }
+
+ // buffer handles metrics
+ if (!buffer.offer(record)) {
+ if (LOGGER.isDebugEnabled()) {
+ LOGGER.debug(
+ "publish() - buffer.offer() rejected, dropping record:{}", record.getMessage());
+ }
+ } else if (LOGGER.isTraceEnabled()) {
+ LOGGER.trace("publish() - buffer.offer() accepted, record:{}", record.getMessage());
+ }
+ }
+
+ /**
+ * Wakes up the uploading thread to check the buffer for batches.
+ *
+ *
This will only drain the buffer fully (i.e. including partial batches) if {@link #close()}
+ * is called or {@link #isClosed} is set.
+ */
+ @Override
+ public void flush() {
+ maybeTrace("flush() - called");
+ notifyUploading();
+ }
+
+ /**
+ * Closes the LogHandler so that it will drop any records sent to {@link #publish(LogRecord)} and
+ * drain the buffer fully to send all batches to S3.
+ */
+ @Override
+ public void close() {
+
+ LOGGER.info(
+ "closing() - marking handler closed, flushing, and waiting for uploads to complete. uploadShutdownDeadline:{}",
+ uploadShutdownDeadline);
+
+ // mark as closed to stop accepting further events and tell the upload thread
+ // to drain the buffer fully.
+ isClosed.set(true);
+ flush();
+
+ try {
+ // Check uploading is not running by trying to get the single uploading permit
+ // TODO: move timeout to config
+ if (!uploadPermit.tryAcquire(uploadShutdownDeadline.toMillis(), TimeUnit.MILLISECONDS)) {
+ LOGGER.warn(
+ "close() - Failed to get uploading permit, upload failed to stop. uploadShutdownDeadline:{}",
+ uploadShutdownDeadline);
+ } else {
+ uploadPermit.release();
+ LOGGER.debug("close() - acquired uploading permit, uploading has completed.");
+ }
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ LOGGER.warn("close() - Interrupted waiting for billing upload loop to finish");
+ } finally {
+ uploader.close();
+ }
+ }
+
+ // ============================================================
+ // Flush pipeline
+ // ============================================================
+
+ /**
+ * Call this on a worker thread to start uploading, will start a loop of waiting for batches from
+ * the buffer and uploading them.
+ */
+ void startUploading() {
+
+ LOGGER.info("startUploading() - handler:{}, buffer:{}, uploader:{}", this, buffer, uploader);
+
+ boolean hasPermit = false;
+ BatchedLogBuffer.Batch batch;
+ try {
+ if (!(hasPermit = uploadPermit.tryAcquire())) {
+ throw new IllegalStateException(
+ "startUploading() - unable to acquire uploadPermit, was function already called?");
+ }
+
+ while (true) {
+
+ // if the handler is closed we do not want to go to sleep again because it is closing
+ // down.
+ if (!isClosed.get()) {
+ try {
+ // waiting will release the synchronized monitor
+ maybeTrace(
+ "startUploading() - waiting for wakeupPermit. isClosed:{}, uploadSleepDuration:{}",
+ isClosed.get(),
+ uploadSleepDuration);
+ var acquiredWakePermit =
+ wakeupPermit.tryAcquire(uploadSleepDuration.toMillis(), TimeUnit.MILLISECONDS);
+ // is not important if we got a permit to wake, or timed out, just for logging
+ maybeTrace(
+ "startUploading() - wakeup permit or timeout, isClosed:{}, acquiredWakePermit:{}",
+ isClosed.get(),
+ acquiredWakePermit);
+
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ dumpBuffer();
+ return;
+ }
+ } else {
+ maybeTrace(
+ "startUploading() - not waiting for wakeupPermit because isClosed:{}",
+ isClosed.get());
+ }
+
+ // Get the next batches, if isClosed is true then we want to drain all events
+ // which may mean creating a batch when we do not have a full one.
+ while ((batch = buffer.nextBatch(isClosed.get())) != null) {
+ uploadBatch(batch);
+ }
+
+ if (isClosed.get()) {
+ // Handler is closing down, time to get out of this crazy loop
+ break;
+ }
+ }
+ } finally {
+ // release the uploading permit if we have it, done with the uploading lifestyle
+ if (hasPermit) {
+ uploadPermit.release();
+ maybeTrace("startUploading() - releasing upload permit");
+ } else {
+ maybeTrace("startUploading() - upload permit was not acquired, not releasing");
+ }
+ }
+
+ if (!buffer.isEmpty()) {
+ LOGGER.warn(
+ "startUploading() - finished with abandoned billing events, billingQueue.size():{} ",
+ buffer.size());
+ }
+
+ LOGGER.info(
+ "startUploading() - stopped uploading. handler:{}, buffer:{}, uploader:{}",
+ this,
+ buffer,
+ uploader);
+ }
+
+ /** Adds a permit to the wakeupPermit so the uploading thread will wakeup and do some work. */
+ private void notifyUploading() {
+ wakeupPermit.release();
+ }
+
+ private static void maybeTrace(String message, Object... args) {
+ if (LOGGER.isTraceEnabled()) {
+ LOGGER.trace(message, args);
+ }
+ }
+
+ /**
+ * Creates a Uni that will upload the provided batch.
+ *
+ *
As a deferred Uni it does not do any work until something pulls the item, so the caller (see
+ * startUploading()) starts the work and can decide to wait etc.
+ *
+ * @param batch
+ * @return
+ */
+ private void uploadBatch(BatchedLogBuffer.Batch batch) {
+
+ LOGGER.info(
+ "uploadBatch() - starting to upload. uploaderSafetyDeadline:{}, batch:{}",
+ uploaderSafetyDeadline,
+ batch);
+
+ // while the uploader should take of all the timeout and retry logic
+ // as a client of the uploader adding a safety timeout here incase it breaks
+
+ // using deferred so that an error in upload() before it returns the Uni is then
+ // treated as an error through the Uni pipeline
+ var uploadResult =
+ Uni.createFrom()
+ .deferred(() -> uploader.upload(batch))
+ .ifNoItem()
+ .after(uploaderSafetyDeadline)
+ .fail()
+ .onFailure()
+ .recoverWithItem(
+ t -> onUploaderFailure(batch, t)) // TimeoutException id deadline exceeded
+ .await()
+ .indefinitely(); // the deadline above covers it
+
+ if (uploadResult.throwable() == null) {
+ onBatchSuccess(uploadResult);
+ } else {
+ onBatchFailure(uploadResult);
+ }
+ }
+
+ /**
+ * There was an unhandled error from the uploader().
+ *
+ *
Could be from in upload() before it returned or from running the Uni to do the upload. Just
+ * map this unhandled back into the UploadResult so we can deal with error in standard way
+ */
+ private AsyncBatchedLogUploader.UploadResult onUploaderFailure(
+ BatchedLogBuffer.Batch batch, Throwable throwable) {
+ LOGGER.error(
+ "onUploaderFailure() - throwable from uploader, adding to UploadResult. batch:{}, throwable:{}",
+ batch,
+ throwable.toString());
+ return new AsyncBatchedLogUploader.UploadResult(batch, throwable);
+ }
+
+ private void onBatchSuccess(AsyncBatchedLogUploader.UploadResult uploadResult) {
+ LOGGER.info("onBatchSuccess() - successfully uploaded batch:{}", uploadResult.batch());
+ }
+
+ private void onBatchFailure(AsyncBatchedLogUploader.UploadResult uploadResult) {
+ LOGGER.error("onBatchFailure() - failed to upload batch:{}", uploadResult.batch());
+ }
+
+ /** TODO: dump buffer or a failed batch to regular logs or whatever */
+ private void dumpBuffer() {}
+
+ private void dumpBatch() {}
+
+ @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/provider/DefaultBilling.java b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/DefaultBilling.java
similarity index 98%
rename from src/main/java/io/stargate/sgv2/jsonapi/service/provider/DefaultBilling.java
rename to src/main/java/io/stargate/sgv2/jsonapi/service/billing/DefaultBilling.java
index 9ab810bd4a..b62e8c5ffb 100644
--- a/src/main/java/io/stargate/sgv2/jsonapi/service/provider/DefaultBilling.java
+++ b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/DefaultBilling.java
@@ -1,4 +1,4 @@
-package io.stargate.sgv2.jsonapi.service.provider;
+package io.stargate.sgv2.jsonapi.service.billing;
import static io.stargate.sgv2.jsonapi.util.StringUtil.requireNonBlank;
@@ -8,6 +8,7 @@
import com.google.common.annotations.VisibleForTesting;
import io.stargate.sgv2.jsonapi.config.BillingConfig;
import io.stargate.sgv2.jsonapi.config.feature.ApiFeature;
+import io.stargate.sgv2.jsonapi.service.provider.ModelUsage;
import java.time.Instant;
import java.util.ArrayList;
import java.util.List;
diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/billing/S3BatchedLogUploader.java b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/S3BatchedLogUploader.java
new file mode 100644
index 0000000000..5369194ec8
--- /dev/null
+++ b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/S3BatchedLogUploader.java
@@ -0,0 +1,252 @@
+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 io.smallrye.mutiny.infrastructure.Infrastructure;
+import io.stargate.sgv2.jsonapi.metrics.BatchedLogUploaderMetrics;
+import java.net.URI;
+import java.nio.charset.StandardCharsets;
+import java.time.Duration;
+import java.time.ZoneOffset;
+import java.time.format.DateTimeFormatter;
+import java.util.Objects;
+import java.util.concurrent.CompletionException;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import software.amazon.awssdk.awscore.exception.AwsServiceException;
+import software.amazon.awssdk.core.async.AsyncRequestBody;
+import software.amazon.awssdk.core.retry.RetryMode;
+import software.amazon.awssdk.regions.Region;
+import software.amazon.awssdk.services.s3.S3AsyncClient;
+import software.amazon.awssdk.services.s3.model.PutObjectRequest;
+
+/**
+ * Uploads sealed billing batches to S3 as NDJSON objects under time-partitioned keys. TODO:
+ * .requestChecksumCalculation(RequestChecksumCalculation.WHEN_SUPPORTED)
+ * .responseChecksumValidation(ResponseChecksumValidation.WHEN_SUPPORTED)
+ */
+public class S3BatchedLogUploader implements AsyncBatchedLogUploader {
+
+ private static final Logger LOGGER = LoggerFactory.getLogger(S3BatchedLogUploader.class);
+
+ // S3 destination formatting
+ // private static final String PATH_PREFIX = "data-api";
+ // 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 static final Duration API_CALL_ATTEMPT_TIMEOUT = Duration.ofSeconds(10);
+ // private static final Duration API_CALL_TIMEOUT = Duration.ofSeconds(30);
+
+ private final S3AsyncClient client;
+ private final String region;
+ private final String bucket;
+ private final 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
+ 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;
+ }
+
+ /**
+ * Creates a new instance
+ *
+ * @param region
+ * @param bucket
+ * @param endpointOverride
+ * @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://.s3..amazonaws.com from region + partition dnsSuffix.
+ // An override is only for a non-AWS S3 (S3Mock in tests): it bypasses those rules and forces
+ // path-style, since a localhost host can't virtual-host the bucket as a subdomain.
+ if (endpointOverride != null) {
+ LOGGER.warn(
+ "create() - WARNING - using endpointOverride this should only be used in testing. endpointOverride: {}",
+ endpointOverride);
+ builder.endpointOverride(URI.create(endpointOverride)).forcePathStyle(Boolean.TRUE);
+ }
+
+ return new S3BatchedLogUploader(builder.build(), region, bucket, pathPrefix, metrics);
+ }
+
+ @Override
+ public Uni upload(BatchedLogBuffer.Batch batch) {
+
+ Objects.requireNonNull(batch, "batch must not be null");
+
+ var location = objectLocation(batch);
+ var body = objectContent(batch);
+
+ LOGGER.info(
+ "upload() - starting to upload batch, batch:{}, location:{}, body.size:{}",
+ batch,
+ location,
+ body.length);
+
+ // retry and timeout are set when we created the client.
+
+ var putRequest =
+ PutObjectRequest.builder()
+ .bucket(location.bucket())
+ .key(location.key())
+ .contentType(CONTENT_TYPE_NDJSON)
+ .build();
+ if (LOGGER.isDebugEnabled()) {
+ LOGGER.debug("upload() - batch:{}, putRequest: {}", batch, putRequest);
+ }
+
+ // Call to S3 client comes back on its own worker thread
+ // using emit() so following processing happens on a quarkus worker thread
+ return Uni.createFrom()
+ .completionStage(client.putObject(putRequest, AsyncRequestBody.fromBytes(body)))
+ .emitOn(Infrastructure.getDefaultWorkerPool())
+ .onItemOrFailure()
+ .transform(
+ (resp, failure) -> {
+ var success = failure == null;
+ var cause = (failure instanceof CompletionException) ? failure.getCause() : failure;
+ var requestId = (cause instanceof AwsServiceException ase) ? ase.requestId() : null;
+
+ if (!success) {
+ metrics.recordBatchFailed(batch);
+ LOGGER.error(
+ "upload() - error uploading billing to S3. batch:{}, location:{}, requestId:{}",
+ batch,
+ location,
+ requestId,
+ cause);
+ } else {
+ metrics.recordBatchDelivered(batch);
+ LOGGER.info(
+ "upload() - success uploading billing to S3. batch:{}, location:{}, requestId:{}, eTag:{}, status:{}",
+ batch,
+ location,
+ requestId,
+ resp.eTag(),
+ resp.sdkHttpResponse().statusCode());
+ }
+ return new UploadResult(batch, failure);
+ });
+ }
+
+ @Override
+ public void close() {
+ client.close();
+ }
+
+ @Override
+ public String toString() {
+ return new StringBuilder(classSimpleName(this) + "{")
+ .append("region=")
+ .append(region)
+ .append(", bucket=")
+ .append(bucket)
+ .append(", pathPrefix=")
+ .append(pathPrefix)
+ .append("}")
+ .toString();
+ }
+
+ @VisibleForTesting
+ S3Location objectLocation(BatchedLogBuffer.Batch batch) {
+
+ var objectKey =
+ pathPrefix
+ + "/"
+ + OBJECT_KEY_FORMATTER.format(batch.oldestEventAt())
+ + "/"
+ + batch.id()
+ + ".jsonl";
+
+ if (LOGGER.isDebugEnabled()) {
+ LOGGER.debug("objectLocation() - batch:{}, objectKey:{}", batch, objectKey);
+ }
+ return new S3Location(region, bucket, objectKey);
+ }
+
+ @VisibleForTesting
+ byte[] objectContent(BatchedLogBuffer.Batch batch) {
+
+ StringBuilder sb = new StringBuilder();
+ for (String line : batch.lines()) {
+ sb.append(line).append('\n');
+ }
+ var bytes = sb.toString().getBytes(StandardCharsets.UTF_8);
+
+ if (LOGGER.isDebugEnabled()) {
+ LOGGER.debug("objectContent() - batch:{}, bytes.length:{}", batch, bytes.length);
+ }
+ return bytes;
+ }
+
+ @VisibleForTesting
+ record S3Location(String region, String bucket, String key) {}
+}
diff --git a/src/test/java/io/stargate/sgv2/jsonapi/TestConstants.java b/src/test/java/io/stargate/sgv2/jsonapi/TestConstants.java
index e0c42122bc..8dd3949a92 100644
--- a/src/test/java/io/stargate/sgv2/jsonapi/TestConstants.java
+++ b/src/test/java/io/stargate/sgv2/jsonapi/TestConstants.java
@@ -18,11 +18,11 @@
import io.stargate.sgv2.jsonapi.config.constants.DocumentConstants;
import io.stargate.sgv2.jsonapi.config.feature.ApiFeatures;
import io.stargate.sgv2.jsonapi.metrics.JsonProcessingMetricsReporter;
+import io.stargate.sgv2.jsonapi.service.billing.Billing;
import io.stargate.sgv2.jsonapi.service.cqldriver.CQLSessionCache;
import io.stargate.sgv2.jsonapi.service.cqldriver.executor.*;
import io.stargate.sgv2.jsonapi.service.embedding.operation.EmbeddingProvider;
import io.stargate.sgv2.jsonapi.service.embedding.operation.EmbeddingProviderFactory;
-import io.stargate.sgv2.jsonapi.service.provider.Billing;
import io.stargate.sgv2.jsonapi.service.reranking.operation.RerankingProviderFactory;
import io.stargate.sgv2.jsonapi.service.schema.*;
import io.stargate.sgv2.jsonapi.service.schema.collections.CollectionLexicalDefSchemaFactory;
diff --git a/src/test/java/io/stargate/sgv2/jsonapi/api/v1/BillingS3ExportIntegrationTest.java b/src/test/java/io/stargate/sgv2/jsonapi/api/v1/BillingS3ExportIntegrationTest.java
new file mode 100644
index 0000000000..b121a50540
--- /dev/null
+++ b/src/test/java/io/stargate/sgv2/jsonapi/api/v1/BillingS3ExportIntegrationTest.java
@@ -0,0 +1,265 @@
+package io.stargate.sgv2.jsonapi.api.v1;
+
+import static io.restassured.RestAssured.given;
+import static io.stargate.sgv2.jsonapi.api.v1.ResponseAssertions.responseIsDDLSuccess;
+import static io.stargate.sgv2.jsonapi.api.v1.ResponseAssertions.responseIsWriteSuccess;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.awaitility.Awaitility.await;
+import static org.hamcrest.Matchers.is;
+
+import com.fasterxml.jackson.databind.JsonNode;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import io.quarkus.test.common.WithTestResource;
+import io.quarkus.test.junit.QuarkusIntegrationTest;
+import io.stargate.sgv2.jsonapi.testresource.DseTestResource;
+import io.stargate.sgv2.jsonapi.testresource.S3MockTestResource;
+import java.net.URI;
+import java.time.Duration;
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Set;
+import java.util.regex.Pattern;
+import org.junit.jupiter.api.MethodOrderer;
+import org.junit.jupiter.api.Order;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.TestMethodOrder;
+import software.amazon.awssdk.auth.credentials.AwsBasicCredentials;
+import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider;
+import software.amazon.awssdk.regions.Region;
+import software.amazon.awssdk.services.s3.S3Client;
+import software.amazon.awssdk.services.s3.model.S3Object;
+
+/**
+ * End-to-end test of the billing S3 export: real vectorize commands (via {@code
+ * CustomITEmbeddingProvider}) emit {@code billing.events} lines, and the installed {@code
+ * BillingS3LogHandler} must land them in the S3Mock bucket as time-partitioned NDJSON objects.
+ *
+ * {@link S3MockTestResource} enables the export with small thresholds (count seal 5, age sweep
+ * 2s) and turns on the {@code billing-events-logging} feature flag.
+ *
+ *
Methods are ordered: the last test stops the S3Mock container to prove a failing export never
+ * affects the data API, which kills S3 for the rest of the class — nothing may run after it.
+ */
+@QuarkusIntegrationTest
+@WithTestResource(value = DseTestResource.class)
+@WithTestResource(value = S3MockTestResource.class)
+@TestMethodOrder(MethodOrderer.OrderAnnotation.class)
+public class BillingS3ExportIntegrationTest extends AbstractKeyspaceIntegrationTestBase {
+
+ private static final String COLLECTION = "billing_export_collection";
+
+ /** Every vectorize call emits at least one billing event, so lines >= documents. */
+ private static final int DOCUMENTS = 10;
+
+ private static final Pattern KEY_PATTERN =
+ Pattern.compile("data-api/\\d{4}/\\d{2}/\\d{2}/\\d{2}/\\d{2}/[0-9a-f-]{36}\\.jsonl");
+
+ /** Wire contract of {@code BillingEventType}: billing consumers key on these exact values. */
+ private static final Set EVENT_TYPES =
+ Set.of(
+ "internal_model_total_tokens",
+ "external_model_total_tokens",
+ "internal_model_egress_bytes",
+ "external_model_egress_bytes",
+ "internal_model_ingress_bytes",
+ "external_model_ingress_bytes");
+
+ private static final ObjectMapper MAPPER = new ObjectMapper();
+
+ @Test
+ public void billingEventsLandInS3AsNdjson() throws Exception {
+ createVectorizeCollection();
+ for (int i = 0; i < DOCUMENTS; i++) {
+ insertDocumentWithVectorize(i);
+ }
+
+ try (S3Client s3 = verificationClient()) {
+ // The count seal ships full batches immediately; the 2s age tick sweeps the remainder.
+ await()
+ .atMost(Duration.ofSeconds(60))
+ .pollInterval(Duration.ofSeconds(2))
+ .untilAsserted(
+ () -> assertThat(exportedLines(s3)).hasSizeGreaterThanOrEqualTo(DOCUMENTS));
+
+ // Object layout: time-partitioned keys and NDJSON content type.
+ List objects = exportObjects(s3);
+ assertThat(objects).isNotEmpty();
+ for (S3Object object : objects) {
+ assertThat(object.key()).matches(KEY_PATTERN);
+ }
+ var head = s3.headObject(b -> b.bucket(S3MockTestResource.BUCKET).key(objects.get(0).key()));
+ assertThat(head.contentType()).isEqualTo("application/x-ndjson");
+
+ // Every line is a self-contained billing event with the expected shape; ids never repeat
+ // across objects. (region/resource_id may be absent locally and are not asserted.)
+ List lines = exportedLines(s3);
+ Set seenIds = new HashSet<>();
+ for (String line : lines) {
+ JsonNode event = MAPPER.readTree(line);
+ String id = event.path("id").asText();
+ assertThat(id).isNotBlank();
+ assertThat(seenIds.add(id))
+ .as("billing event id duplicated across export: %s", id)
+ .isTrue();
+ assertThat(event.path("timestamp").asText()).isNotBlank();
+ assertThat(event.path("product").asText()).isEqualTo("serverless");
+ assertThat(event.path("event_type").asText()).isIn(EVENT_TYPES);
+ JsonNode properties = event.path("properties");
+ assertThat(properties.path("usage").isIntegralNumber()).isTrue();
+ assertThat(properties.path("usage").asLong()).isGreaterThanOrEqualTo(0L);
+ assertThat(properties.path("resource_type").asText()).isEqualTo("serverless_database");
+ assertThat(properties.path("provider").asText()).isEqualTo("custom");
+ // The billed model is what the provider reports in ModelUsage — for the IT provider that
+ // is its internal model config ("test-model"), not the createCollection modelName.
+ assertThat(properties.path("model").asText()).isEqualTo("test-model");
+ }
+ }
+
+ // The delivery counters on /metrics must agree that the export is alive.
+ await()
+ .atMost(Duration.ofSeconds(10))
+ .untilAsserted(
+ () ->
+ assertThat(metricTotal("billing_s3_events_flushed_total"))
+ .isGreaterThanOrEqualTo(DOCUMENTS));
+ }
+
+ /**
+ * Billing is a side-channel: a failing S3 export must never affect the data API. Stopping the
+ * S3Mock container leaves the endpoint dead — every upload from here on fails with
+ * connection-refused, like an S3 outage — yet inserts must keep returning normal write successes,
+ * and the failures must be counted rather than silently swallowed. The handler's failure
+ * accounting in isolation is covered by {@code BillingS3LogHandlerTest}; this proves the property
+ * end-to-end in the packaged app.
+ *
+ * Must run last ({@link S3MockTestResource#stopContainer()} is one-way): any test needing a
+ * live S3 goes before this one. Reuses the collection created by the happy-path test.
+ *
+ *
{@code Integer.MAX_VALUE}, not a small sentinel, and deliberately the only ordered method:
+ * {@code OrderAnnotation} gives an unannotated method the default order {@code Integer.MAX_VALUE
+ * / 2}, so any newly added test with no {@code @Order} still sorts before this one. Do NOT lower
+ * this value — anything below the default would let such a test run after S3 is dead.
+ */
+ @Test
+ @Order(Integer.MAX_VALUE)
+ public void exportFailureDoesNotAffectTheApi() {
+ S3MockTestResource.stopContainer();
+
+ // Each insert emits billing events whose upload will fail — yet every insert must still
+ // return a normal write success, because publish() is fire-and-forget and never waits on S3.
+ for (int i = 0; i < DOCUMENTS; i++) {
+ insertDocumentWithVectorize(DOCUMENTS + i);
+ }
+
+ // Failures are counted, not silently swallowed. Uploads settle as failed only after the SDK
+ // exhausts its retries, so poll for the counter to move.
+ await()
+ .atMost(Duration.ofSeconds(60))
+ .pollInterval(Duration.ofSeconds(2))
+ .untilAsserted(
+ () -> assertThat(metricTotal("billing_s3_batches_failed_total")).isGreaterThan(0.0));
+
+ // The API is still healthy after the export has been failing for a while: one more insert
+ // succeeds exactly like the first.
+ insertDocumentWithVectorize(2 * DOCUMENTS);
+ }
+
+ // ============================================================
+ // Command helpers
+ // ============================================================
+
+ private void createVectorizeCollection() {
+ givenHeadersPostJsonThenOk(
+ """
+ {
+ "createCollection": {
+ "name": "%s",
+ "options": {
+ "vector": {
+ "metric": "cosine",
+ "dimension": 5,
+ "service": {
+ "provider": "custom",
+ "modelName": "text-embedding-ada-002",
+ "authentication": {
+ "providerKey" : "shared_creds.providerKey"
+ },
+ "parameters": {
+ "projectId": "test project"
+ }
+ }
+ }
+ }
+ }
+ }
+ """
+ .formatted(COLLECTION))
+ .body("$", responseIsDDLSuccess())
+ .body("status.ok", is(1));
+ }
+
+ private void insertDocumentWithVectorize(int i) {
+ String json =
+ """
+ {
+ "insertOne": {
+ "document": {
+ "_id": "doc-%d",
+ "description": "billing export test document %d",
+ "$vectorize": "billing export test document %d"
+ }
+ }
+ }
+ """
+ .formatted(i, i, i);
+ givenHeadersAndJson(json)
+ .when()
+ .post(CollectionResource.BASE_PATH, keyspaceName, COLLECTION)
+ .then()
+ .statusCode(200)
+ .body("$", responseIsWriteSuccess());
+ }
+
+ // ============================================================
+ // S3 verification helpers
+ // ============================================================
+
+ private static S3Client verificationClient() {
+ return S3Client.builder()
+ .region(Region.of(S3MockTestResource.BUCKET_REGION))
+ .credentialsProvider(
+ StaticCredentialsProvider.create(
+ AwsBasicCredentials.create(
+ S3MockTestResource.ACCESS_KEY, S3MockTestResource.SECRET_KEY)))
+ .endpointOverride(URI.create(S3MockTestResource.endpoint()))
+ .forcePathStyle(true)
+ .build();
+ }
+
+ private static List exportObjects(S3Client s3) {
+ return s3.listObjectsV2(b -> b.bucket(S3MockTestResource.BUCKET).prefix("data-api/"))
+ .contents();
+ }
+
+ private static List exportedLines(S3Client s3) {
+ List lines = new ArrayList<>();
+ for (S3Object object : exportObjects(s3)) {
+ String body =
+ s3.getObjectAsBytes(b -> b.bucket(S3MockTestResource.BUCKET).key(object.key()))
+ .asUtf8String();
+ body.lines().filter(line -> !line.isBlank()).forEach(lines::add);
+ }
+ return lines;
+ }
+
+ /** Sum of one counter across all tag combinations on {@code /metrics} (0 when absent). */
+ private static double metricTotal(String metricName) {
+ String metrics = given().when().get("/metrics").then().statusCode(200).extract().asString();
+ return metrics
+ .lines()
+ .filter(line -> line.startsWith(metricName))
+ .mapToDouble(line -> Double.parseDouble(line.substring(line.lastIndexOf(' ') + 1)))
+ .sum();
+ }
+}
diff --git a/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BatchedLogBufferTest.java b/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BatchedLogBufferTest.java
new file mode 100644
index 0000000000..4af4868e34
--- /dev/null
+++ b/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BatchedLogBufferTest.java
@@ -0,0 +1,594 @@
+package io.stargate.sgv2.jsonapi.service.billing;
+
+import static io.stargate.sgv2.jsonapi.util.ClassUtils.classSimpleName;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+import static org.awaitility.Awaitility.await;
+import static org.mockito.Mockito.*;
+
+import io.stargate.sgv2.jsonapi.metrics.BatchedLogBufferMetrics;
+import java.lang.ref.WeakReference;
+import java.time.Duration;
+import java.util.*;
+import java.util.concurrent.*;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.logging.Level;
+import java.util.logging.LogRecord;
+import java.util.stream.Stream;
+import org.junit.jupiter.api.Test;
+
+/**
+ * Unit tests for {@link BatchedLogBuffer}\
+ *
+ * TODO: out of order log records gets correct oldest metric TODO: TEST a big line bigger than
+ * the max bytes gets through TODO: test metrics using SimpleMeterRegistry
+ */
+public class BatchedLogBufferTest 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 snapshot = BufferSnapshot.create(fixture);
+ var slice = Slice.to(BUFFER_CAPACITY);
+
+ // send full capacity to the buffer, should all work
+ fixture.assertOffer("offerFailsAtCapacitySingleThread() - prefill to capacity", slice);
+
+ // check the change in the buffer is expected given the slice of source data
+ snapshot.assertAll("offerFailsAtCapacitySingleThread()", slice, true);
+ // Buffer should now be full, try to add one more
+ fixture.assertBufferFull("offerFailsAtCapacitySingleThread()", BUFFER_CAPACITY + 1);
+ }
+
+ /** When the buffer reaches capacity calling offer() fails. Multiple producer threads. */
+ @Test
+ public void offerFailsAtCapacityMultiThread() {
+
+ var fixture = 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 batchIds = new HashSet<>();
+ while ((batch = fixture.buffer().nextBatch(true)) != null) {
+
+ assertThat(batch.id())
+ .as("nextBatchBatchProperties() - batch ID has not been seen")
+ .satisfies(batchIds::add);
+
+ assertThat(batch.toString())
+ .as("nextBatchBatchProperties() - batch toString has values")
+ .contains("id=" + batch.id())
+ .contains("reason=" + batch.reason())
+ .contains("size=" + batch.size())
+ .contains("bytes=" + batch.bytes());
+ }
+
+ assertThat(fixture.buffer().isEmpty())
+ .as("nextBatchBatchProperties() - drained buffer is empty")
+ .isTrue();
+ }
+
+ /** Metadata (size etc) for the buffer is updated after a batch is returned. */
+ @Test
+ public void nextBatchMetaUpdatedAfterBatch() {
+
+ // lock the clock, do not want it to auto advance for batch testing
+ var fixture = defaultBufferFixture(true);
+ var slice = Slice.to(MAX_BATCH_SIZE);
+
+ // fill the buffer with 1 batch size and assert metadata
+ fixture.assertOffer("nextBatchMetaUpdatedAfterBatch()", slice);
+
+ // take 1 batch
+ // we are only checking that the bookkeeping on the buffer changes, not checking
+ // rules for batch selections, this is done by assertNextBatch()
+ var batch1 = fixture.assertNextBatch("nextBatchMetaUpdatedAfterBatch() - 1st", false);
+
+ // take a second batch and check again bookkeeping updated
+ var batch2 = fixture.assertNextBatch("nextBatchMetaUpdatedAfterBatch() - 2nd", false);
+ }
+
+ /** Trigger a batch from the number of records added to buffer */
+ @Test
+ public void nextBatchTriggerMaxSize() {
+
+ // change so the template is small so does not trigger max bytes
+ // lock the clock, do not want it to auto advance for batch testing
+ var fixture =
+ Fixture.createFixture(
+ MAX_BATCH_SIZE,
+ MAX_BATCH_BYTES * 100, // big number so never batch because of bytes
+ MAX_AGE,
+ BUFFER_CAPACITY,
+ NUM_RECORDS,
+ LOG_LEVEL,
+ "test-",
+ true,
+ false,
+ false,
+ false);
+
+ // Fill to 1 less than max batch size, should be no batch
+ var slice1 = Slice.to(MAX_BATCH_SIZE - 1);
+ fixture.assertOffer("nextBatchMetaUpdatedAfterBatch()", slice1);
+ var batch1 = fixture.buffer().nextBatch(false);
+ assertThat(batch1).as("nextBatchTriggerMaxSize() - < MAX_BATCH_SIZE, no batch").isNull();
+
+ // add one more record, should be a batch of MAX_BATCH_SIZE
+ var slice2 = Slice.slice(MAX_BATCH_SIZE - 1, MAX_BATCH_SIZE);
+ fixture.assertOffer("nextBatchMetaUpdatedAfterBatch()", slice2);
+ var batch2 = fixture.assertNextBatch("nextBatchMetaUpdatedAfterBatch() - 2nd", false);
+
+ assertThat(batch2.size())
+ .as("nextBatchTriggerMaxSize() - 2nd batch is full batch size")
+ .isEqualTo(MAX_BATCH_SIZE);
+
+ assertThat(batch2.reason())
+ .as(
+ "nextBatchTriggerMaxSize() - 2nd batch because "
+ + BatchedLogBuffer.BillingBatchReason.MAX_SIZE_EXCEEDED)
+ .isEqualTo(BatchedLogBuffer.BillingBatchReason.MAX_SIZE_EXCEEDED);
+
+ // add one more record, should be no more batches
+ var slice3 = Slice.slice(MAX_BATCH_SIZE, MAX_BATCH_SIZE + 1);
+ fixture.assertOffer("nextBatchMetaUpdatedAfterBatch() - 3rd", slice3);
+ var batch3 = fixture.buffer().nextBatch(false);
+ assertThat(batch3).as("nextBatchTriggerMaxSize() - 3rd - no batch").isNull();
+ }
+
+ /** Trigger a batch from the byte size in the buffer */
+ @Test
+ public void nextBatchTriggerMaxBytes() {
+
+ // default fixture will only fit
+ // the MAX_BATCH_BYTES_NUM_MESSAGES which is less than MAX_SIZE
+ // lock the clock, do not want it to auto advance for batch testing
+ var fixture = defaultBufferFixture(true);
+
+ // Fill to 1 message less than max bytes size, should be no batch
+ var slice1 = Slice.to(MAX_BATCH_BYTES_NUM_MESSAGES - 1);
+ fixture.assertOffer("nextBatchTriggerMaxBytes()", slice1);
+ var batch1 = fixture.buffer().nextBatch(false);
+ assertThat(batch1).as("nextBatchTriggerMaxSize() - < MAX_BATCH_BYTES, no batch").isNull();
+
+ // add one more , should be a batch of full batch bytes
+ var slice2 = Slice.slice(MAX_BATCH_BYTES_NUM_MESSAGES - 1, MAX_BATCH_BYTES_NUM_MESSAGES);
+ fixture.assertOffer("nextBatchTriggerMaxBytes()", slice2);
+ var batch2 = fixture.assertNextBatch("nextBatchTriggerMaxBytes() - 2nd", false);
+
+ // we know how many we put in there
+ assertThat(batch2.bytes())
+ .as("nextBatchTriggerMaxBytes() - 2nd batch byte size match")
+ .isEqualTo(MAX_BATCH_BYTES_NUM_MESSAGES * MESSAGE_LENGTH_IN_BUFFER);
+
+ assertThat(batch2.reason())
+ .as(
+ "nextBatchTriggerMaxBytes() - 2nd batch because "
+ + BatchedLogBuffer.BillingBatchReason.MAX_BYTES_EXCEEDED)
+ .isEqualTo(BatchedLogBuffer.BillingBatchReason.MAX_BYTES_EXCEEDED);
+
+ // add one more, should be no more batches
+ var slice3 = Slice.slice(MAX_BATCH_BYTES_NUM_MESSAGES, MAX_BATCH_BYTES_NUM_MESSAGES + 1);
+ fixture.assertOffer("nextBatchTriggerMaxBytes() - 3rd", slice3);
+ var batch3 = fixture.buffer().nextBatch(false);
+ assertThat(batch3).as("nextBatchTriggerMaxBytes() - 3rd - no batch").isNull();
+ }
+
+ /** Trigger a batch from the maximum age of the first element in the buffer */
+ @Test
+ public void nextBatchTriggerMaxAge() {
+
+ // lock the clock, do not want it to auto advance for batch testing
+ // NOTE: WE ARE USING THE MOCK CLOCK IN THIS TEST, WE CONTROL TIME
+ var fixture = defaultBufferFixture(true);
+
+ // Add only 3 messages, we will not trip size or bytes tigger
+ final int ADDED_RECORDS = 3;
+ var slice1 = Slice.to(ADDED_RECORDS);
+ fixture.assertOffer("nextBatchTriggerMaxAge()", slice1);
+
+ // the clock has not moved, there should be no batch
+ var batch1 = fixture.buffer().nextBatch(false);
+ assertThat(batch1).as("nextBatchTriggerMaxAge() - clock as not moved, no batch").isNull();
+
+ // Every LogRecord created in fixture has an instanceAt of 1 second after the previous
+ // the first LogRecord has the same instanceAt as when the clock started.
+ // so if we advance the clock to be MAX_AGE after when it started the only LogRecord that will
+ // be too old is the first, the others are all 1+ seconds younger
+ var newNow = fixture.clock().startedAt().plus(MAX_AGE);
+ fixture.clock().setInstant(newNow);
+
+ // The buffer should now think the time is "newNow"
+ // Sanity check, before getting the batch check that only the first log record is MAX_AGE
+ // checking all this junk did what I think
+ int i = 0;
+ var peekedBuffer = fixture.buffer().peekBuffer();
+ for (var peekEntry : peekedBuffer) {
+ var entryAge = fixture.buffer().entryAge(peekEntry);
+ if (i == 0) {
+ assertThat(entryAge)
+ .as("nextBatchTriggerMaxAge() - clock moved, first entry should be MAX_AGE old")
+ .isEqualTo(MAX_AGE);
+ } else {
+ assertThat(entryAge)
+ .as(
+ "nextBatchTriggerMaxAge() - clock moved, non first entry should be < MAX_AGE old. i: "
+ + i)
+ .isLessThan(MAX_AGE);
+ }
+ i++;
+ }
+
+ // with the clock advanced the buffer should now trigger a batch because
+ // MAX_AGE_EXCEEDED
+ var batch2 = fixture.assertNextBatch("nextBatchTriggerMaxAge() - 2nd", false);
+
+ assertThat(batch2.reason())
+ .as(
+ "nextBatchTriggerMaxAge() - 2nd batch because "
+ + BatchedLogBuffer.BillingBatchReason.MAX_AGE_EXCEEDED)
+ .isEqualTo(BatchedLogBuffer.BillingBatchReason.MAX_AGE_EXCEEDED);
+
+ // should have drained all the messages, even if they were not too old
+ assertThat(fixture.buffer().size())
+ .as("nextBatchTriggerMaxAge() - 2nd batch buffer, size")
+ .isEqualTo(0);
+ assertThat(fixture.buffer().isEmpty())
+ .as("nextBatchTriggerMaxAge() - 2nd batch buffer, isEmpty")
+ .isTrue();
+ assertThat(fixture.buffer().queuedBytes())
+ .as("nextBatchTriggerMaxAge() - 2nd batch buffer, bytes")
+ .isEqualTo(0);
+
+ // sanity check, we should have ADDED_RECORDS entries in the batch
+ // and the oldest should be the first one we created
+ assertThat(batch2.size())
+ .as("nextBatchTriggerMaxAge() - 2nd batch buffer, size expected")
+ .isEqualTo(ADDED_RECORDS);
+ assertThat(batch2.oldestEventAt())
+ .as("nextBatchTriggerMaxAge() - 2nd batch buffer, oldest event expected")
+ .isEqualTo(fixture.logRecords().getFirst().getInstant());
+
+ // add one more record, should be no more batches
+ var slice3 = Slice.slice(ADDED_RECORDS, ADDED_RECORDS + 1);
+ fixture.assertOffer("nextBatchTriggerMaxAge() - 3rd", slice3);
+ var batch3 = fixture.buffer().nextBatch(false);
+ assertThat(batch3).as("nextBatchTriggerMaxAge() - 3rd - no batch").isNull();
+ }
+
+ /**
+ * Trigger a batch because drainFully=true so we want everything from it regardless of size,
+ * bytes, age
+ */
+ @Test
+ public void nextBatchTriggerDrain() {
+
+ // lock the clock, do not want it to auto advance for batch testing
+ var fixture = defaultBufferFixture(true);
+
+ // Fill so we have 1 full batch and 1 partial batch
+ var PARTIAL_BATCH_SIZE = 10;
+ var slice1 = Slice.to(MAX_BATCH_BYTES_NUM_MESSAGES + PARTIAL_BATCH_SIZE);
+ fixture.assertOffer("nextBatchTriggerDrain()", slice1);
+
+ // 1st - drainFully - should get a full batch
+ var batch1 = fixture.assertNextBatch("nextBatchTriggerDrain() - 1st - full batch", true);
+ assertThat(batch1.reason())
+ .as(
+ "nextBatchTriggerDrain() - 1st - reason is "
+ + BatchedLogBuffer.BillingBatchReason.DRAINING)
+ .isEqualTo(BatchedLogBuffer.BillingBatchReason.DRAINING);
+ assertThat(batch1.size())
+ .as("nextBatchTriggerDrain() - 1st - full batch, size")
+ .isEqualTo(MAX_BATCH_BYTES_NUM_MESSAGES);
+
+ // 2nd - drainFully - should get a partial batch
+ var batch2 = fixture.assertNextBatch("nextBatchTriggerDrain() - 2nd - partial batch", true);
+ assertThat(batch2.reason())
+ .as(
+ "nextBatchTriggerDrain() - 2nd - reason is "
+ + BatchedLogBuffer.BillingBatchReason.DRAINING)
+ .isEqualTo(BatchedLogBuffer.BillingBatchReason.DRAINING);
+ assertThat(batch2.size())
+ .as("nextBatchTriggerDrain() - 2nd - partial batch, size")
+ .isEqualTo(PARTIAL_BATCH_SIZE);
+
+ // 3rs - drainFully - no more batch
+ var batch3 = fixture.buffer().nextBatch(true);
+ assertThat(batch3).as("nextBatchTriggerMaxBytes() - 3rd - no batch").isNull();
+ }
+
+ /**
+ * Multiple producers sending to the buffer, and one consumer reading from it concurrently.
+ *
+ * NOTE: ttest takes 7 or 8 seconds, if you change the sleep time it may mean there are
+ * no batches collected after shutdown because producers go fast
+ */
+ @Test
+ public void multiThreadedProducerConsumer() {
+
+ var fixture = defaultBufferFixture(false);
+
+ // Setup a Consumer thread, it will keep running until we set consumerShutdown
+ var normalBatches = new ArrayList();
+ var shutdownBatches = new ArrayList();
+ var consumerShutdown = new AtomicBoolean(false);
+ var consumerExecutor =
+ Executors.newSingleThreadExecutor(Thread.ofPlatform().name("consumer-", 0).factory());
+
+ var consumerFuture =
+ consumerExecutor.submit(
+ () -> {
+ // this is the consumer in normal operations, read batches, if none sleep, read again
+ while (!consumerShutdown.get()) {
+ BatchedLogBuffer.Batch consumerNormalBatch;
+ // drainFully=false - because not trying to shutdown
+ while ((consumerNormalBatch = fixture.buffer().nextBatch(false)) != null) {
+ normalBatches.add(consumerNormalBatch);
+ // fake that we do some work with the batch, e.g. upload it
+ threadSleep(50);
+ }
+ // fake the sleep between waking up to check for a batch
+ threadSleep(50);
+ }
+
+ // now into the shutdown mode, so drainFully=true to empty the buffer
+ BatchedLogBuffer.Batch consumerShutdownBatch;
+ while ((consumerShutdownBatch = fixture.buffer().nextBatch(true)) != null) {
+ shutdownBatches.add(consumerShutdownBatch);
+ // fake that we do some work with the batch, e.g. upload it
+ threadSleep(50);
+ }
+ });
+ // the consumer is running async looping waiting for batches from the buffer
+
+ // Now setup producers to send data for it, we are going to send all the records we
+ // created, this will be more than the buffer capacity.
+ var NUM_PRODUCER_THREADS = 4;
+ var slice = Slice.to(NUM_RECORDS);
+ var threadFactory = Thread.ofPlatform().name("producer-", 0).factory();
+ var producedCount = new AtomicLong();
+ // we want to pause all producers half way through producing so we can
+ // shutdown the consumer and then produce the remaining records
+ var producerHalfwayLatch = new CountDownLatch(NUM_PRODUCER_THREADS);
+
+ try (var pool = Executors.newFixedThreadPool(NUM_PRODUCER_THREADS, threadFactory)) {
+ for (var record : slice.stream(fixture.logRecords()).toList()) {
+
+ // Append the thread name to the log record for debugging
+ // this will break the config at top of class about how many messages per batch
+ pool.submit(
+ () -> {
+ record.setMessage(
+ record.getMessage() + " - THREAD " + Thread.currentThread().getName());
+ fixture.buffer().offer(record);
+
+ if ((producedCount.incrementAndGet() >= (slice.size() / 2))
+ && (!consumerShutdown.get())) {
+ // this thread got at least half way, mark that and wait for all others
+ // to get this far
+ producerHalfwayLatch.countDown();
+ waitOnLatch(producerHalfwayLatch);
+
+ // Signal the consumer to shut down, next time it wakes it will start using
+ // drainFully
+ consumerShutdown.set(true);
+ } else {
+ // Fake that we are doing other things, do not do if we paused cause we want to
+ // get back to producing ASAP
+ threadSleep(25);
+ }
+ });
+ }
+ } // try, will block waiting for threads to finish when closing Executor
+
+ // wait for consumer to finish
+ try {
+ consumerFuture.get(5, TimeUnit.SECONDS);
+ } catch (InterruptedException | ExecutionException | TimeoutException e) {
+ throw new RuntimeException(e);
+ } finally {
+ // close consumer thread pool
+ consumerExecutor.close();
+ }
+
+ // Now we can check consumer got all the data
+ // all the batches from normal processing should be either max size or bytes
+ for (var batch : normalBatches) {
+
+ assertThat(batch.reason())
+ .as("multiThreadedProducerConsumer() - normal batch reason is size or bytes")
+ .isIn(
+ List.of(
+ BatchedLogBuffer.BillingBatchReason.MAX_SIZE_EXCEEDED,
+ BatchedLogBuffer.BillingBatchReason.MAX_BYTES_EXCEEDED));
+ // sanity check that the messages in the batch came from a producer thread.
+ for (var line : batch.lines()) {
+ assertThat(line)
+ .as("multiThreadedProducerConsumer() - normal batch line created by producer thread.")
+ .contains("THREAD producer-");
+ }
+ }
+
+ // all the batches from shutdown processing must be due to draining
+ for (var batch : shutdownBatches) {
+
+ assertThat(batch.reason())
+ .as("multiThreadedProducerConsumer() - normal batch reason is draining")
+ .isEqualTo(BatchedLogBuffer.BillingBatchReason.DRAINING);
+ // sanity check that the messages in the batch came from a producer thread.
+ for (var line : batch.lines()) {
+ assertThat(line)
+ .as("multiThreadedProducerConsumer() - shutdown batch line created by producer thread.")
+ .contains("THREAD producer-");
+ }
+ }
+
+ // total lines from normal and shutdown must be total from producers
+ var totalBatchLines =
+ Stream.concat(shutdownBatches.stream(), normalBatches.stream())
+ .mapToInt(BatchedLogBuffer.Batch::size)
+ .sum();
+ assertThat(totalBatchLines)
+ .as(
+ "multiThreadedProducerConsumer() - lines from batches same number as produced: "
+ + producedCount.get())
+ .isEqualTo(producedCount.get());
+
+ // sanity check - did every producer thread produce at least one log record ?
+ // log message will look like: "Total of 25 chars 059 - THREAD producer-1"
+ for (int i = 0; i < NUM_PRODUCER_THREADS; i++) {
+ var threadSuffix = "- THREAD producer-" + i;
+ var found =
+ Stream.concat(shutdownBatches.stream(), normalBatches.stream())
+ .flatMap(batch -> batch.lines().stream())
+ .anyMatch(line -> line.endsWith(threadSuffix));
+ assertThat(found).as("Producer thread created record, thread:" + threadSuffix).isTrue();
+ }
+ }
+
+ // *********************************************************
+ // Basic object testing
+ // *********************************************************
+
+ @Test
+ public void testConstructor() {
+
+ var metrics = mock(BatchedLogBufferMetrics.class);
+ assertThatThrownBy(
+ () -> new BatchedLogBuffer(0, 1, Duration.ofSeconds(1), 10, metrics),
+ "maxBatchSize < 1")
+ .isInstanceOf(IllegalArgumentException.class);
+ assertThatThrownBy(
+ () -> new BatchedLogBuffer(1, 0, Duration.ofSeconds(1), 10, metrics),
+ "maxBatchBytes < 1")
+ .isInstanceOf(IllegalArgumentException.class);
+ assertThatThrownBy(
+ () -> new BatchedLogBuffer(1, 1, Duration.ofSeconds(-1), 10, metrics),
+ "maxBatchAge < 1")
+ .isInstanceOf(IllegalArgumentException.class);
+ assertThatThrownBy(
+ () -> new BatchedLogBuffer(1, 1, Duration.ofSeconds(0), 10, metrics), "maxBatchAge = 0")
+ .isInstanceOf(IllegalArgumentException.class);
+ assertThatThrownBy(
+ () -> new BatchedLogBuffer(1, 1, Duration.ofSeconds(1), 0, metrics), "queueCapacity =0")
+ .isInstanceOf(IllegalArgumentException.class);
+
+ clearInvocations(metrics);
+ var buffer = new BatchedLogBuffer(1, 2, Duration.ofSeconds(1), 10, metrics);
+ verify(metrics, times(1).description("buffer registers with metrics")).registerBuffer(any());
+
+ assertThat(buffer.toString())
+ .as("buffer toString has correct values")
+ .startsWith(classSimpleName(buffer))
+ .contains("maxBatchSize=1")
+ .contains("maxBatchBytes=2")
+ .contains("maxBatchAge=PT1S")
+ .contains("size=0");
+ }
+}
diff --git a/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BatchedLogUploaderMetricsTest.java b/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BatchedLogUploaderMetricsTest.java
new file mode 100644
index 0000000000..fab83d80ef
--- /dev/null
+++ b/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BatchedLogUploaderMetricsTest.java
@@ -0,0 +1,51 @@
+package io.stargate.sgv2.jsonapi.service.billing;
+
+/** Guards the meter names and tags — dashboards and alerts key on these exact series. */
+class BatchedLogUploaderMetricsTest {
+ //
+ // @Test
+ // void countersFlowToTheExpectedSeries() {
+ // var registry = new SimpleMeterRegistry();
+ // var metrics = new BatchedLogUploaderMetrics(registry, () -> 0, 100);
+ //
+ // metrics.recordOffered();
+ // metrics.recordDropped();
+ // metrics.recordAbandonedAtShutdown(3);
+ // metrics.recordBatchDelivered(2);
+ // metrics.recordBatchFailed(5);
+ //
+ // assertThat(registry.counter("billing.s3.events.offered").count()).isEqualTo(1.0);
+ // assertThat(registry.counter("billing.s3.events.dropped", "reason", "capacity").count())
+ // .isEqualTo(1.0);
+ // assertThat(registry.counter("billing.s3.events.dropped", "reason", "shutdown").count())
+ // .isEqualTo(3.0);
+ // assertThat(registry.counter("billing.s3.events.flushed").count()).isEqualTo(2.0);
+ // assertThat(registry.counter("billing.s3.batches.uploaded").count()).isEqualTo(1.0);
+ // assertThat(registry.counter("billing.s3.events.failed").count()).isEqualTo(5.0);
+ // assertThat(registry.counter("billing.s3.batches.failed").count()).isEqualTo(1.0);
+ // }
+ //
+ // @Test
+ // void depthGaugeReadsTheLiveSupplier() {
+ // var registry = new SimpleMeterRegistry();
+ // var depth = new AtomicInteger(7);
+ // new BatchedLogUploaderMetrics(registry, depth::get, 100);
+ //
+ // assertThat(registry.get("billing.s3.queue.depth").gauge().value()).isEqualTo(7.0);
+ // depth.set(11);
+ // assertThat(registry.get("billing.s3.queue.depth").gauge().value()).isEqualTo(11.0);
+ // }
+ //
+ // @Test
+ // void deliveryHeartbeatAdvancesOnDeliveredBatches() {
+ // var registry = new SimpleMeterRegistry();
+ // var metrics = new BatchedLogUploaderMetrics(registry, () -> 0, 100);
+ // var heartbeat = registry.get("billing.s3.last_delivery.epoch_seconds").gauge();
+ //
+ // assertThat(heartbeat.value()).isZero(); // never delivered
+ //
+ // long before = Instant.now().getEpochSecond();
+ // metrics.recordBatchDelivered(1);
+ // assertThat(heartbeat.value()).isGreaterThanOrEqualTo(before);
+ // }
+}
diff --git a/src/test/java/io/stargate/sgv2/jsonapi/service/provider/BillingEventTest.java b/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingEventTest.java
similarity index 97%
rename from src/test/java/io/stargate/sgv2/jsonapi/service/provider/BillingEventTest.java
rename to src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingEventTest.java
index dbeca3f937..ee6fbf325e 100644
--- a/src/test/java/io/stargate/sgv2/jsonapi/service/provider/BillingEventTest.java
+++ b/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingEventTest.java
@@ -1,4 +1,4 @@
-package io.stargate.sgv2.jsonapi.service.provider;
+package io.stargate.sgv2.jsonapi.service.billing;
import static net.javacrumbs.jsonunit.JsonAssert.assertJsonEquals;
diff --git a/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3HandlerInstallerTest.java b/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3HandlerInstallerTest.java
new file mode 100644
index 0000000000..18049c3409
--- /dev/null
+++ b/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3HandlerInstallerTest.java
@@ -0,0 +1,165 @@
+package io.stargate.sgv2.jsonapi.service.billing;
+
+import static io.stargate.sgv2.jsonapi.service.billing.BillingS3HandlerInstaller.BILLING_LOGGER_NAME;
+import static org.assertj.core.api.Assertions.assertThat;
+
+import io.micrometer.core.instrument.MeterRegistry;
+import io.micrometer.core.instrument.simple.SimpleMeterRegistry;
+import io.quarkus.runtime.ShutdownEvent;
+import io.quarkus.runtime.StartupEvent;
+import io.smallrye.config.SmallRyeConfigBuilder;
+import io.stargate.sgv2.jsonapi.config.BillingS3ExportConfig;
+import java.util.Arrays;
+import java.util.List;
+import java.util.Map;
+import java.util.logging.Handler;
+import java.util.logging.Logger;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+public class BillingS3HandlerInstallerTest {
+
+ private Logger logger;
+ private Handler[] saved;
+
+ @BeforeEach
+ void saveHandlers() {
+ logger = Logger.getLogger(BILLING_LOGGER_NAME);
+ saved = logger.getHandlers();
+ }
+
+ @AfterEach
+ void restoreHandlers() {
+ for (var h : logger.getHandlers()) {
+ logger.removeHandler(h);
+ }
+ for (var h : saved) {
+ logger.addHandler(h);
+ }
+ }
+
+ @Test
+ void onStartBillingEnabledOthersDisabled() {
+
+ var installer =
+ installer(
+ Map.of(
+ "stargate.jsonapi.billing.s3.enabled",
+ "true",
+ "stargate.jsonapi.billing.s3.disable-other-handlers",
+ "true"),
+ null);
+
+ assertHandlers(installer, true, false);
+ }
+
+ @Test
+ void onStartBillingEnabledOthersEnabled() {
+
+ var installer =
+ installer(
+ Map.of(
+ "stargate.jsonapi.billing.s3.enabled",
+ "true",
+ "stargate.jsonapi.billing.s3.disable-other-handlers",
+ "false"),
+ null);
+
+ assertHandlers(installer, true, true);
+ }
+
+ @Test
+ void onStartBillingDisabledOthersDisabled() {
+
+ var installer =
+ installer(
+ Map.of(
+ "stargate.jsonapi.billing.s3.enabled",
+ "false",
+ "stargate.jsonapi.billing.s3.disable-other-handlers",
+ "true"),
+ null);
+
+ // even though disabling others is enabled, billing s3 is disabled so that should not impact
+ assertHandlers(installer, false, true);
+ }
+
+ @Test
+ void onStartBillingDisabledOthersEnabled() {
+
+ var installer =
+ installer(
+ Map.of(
+ "stargate.jsonapi.billing.s3.enabled",
+ "false",
+ "stargate.jsonapi.billing.s3.disable-other-handlers",
+ "false"),
+ null);
+
+ // even though disabling others is enabled, billing s3 is disabled so that should not impact
+ assertHandlers(installer, false, true);
+ }
+
+ // ====================================
+ // Scaffold
+ // ====================================
+
+ private void assertHandlers(
+ BillingS3HandlerInstaller installer, boolean expectS3Handler, boolean expectOtherHandlers) {
+
+ var billingLogger = Logger.getLogger(BILLING_LOGGER_NAME);
+
+ try {
+ installer.onStart(new StartupEvent());
+ var onStartHandlers = installedHandlers();
+ if (expectS3Handler) {
+ assertThat(onStartHandlers).as("onStart attached S3 handler").contains(installer.handler());
+ }
+ if (expectOtherHandlers) {
+ assertThat(onStartHandlers.size())
+ .as("onStart left other handler in place")
+ .isGreaterThanOrEqualTo(expectS3Handler ? 2 : 1);
+ } else {
+ assertThat(onStartHandlers.size())
+ .as("onStart removed other handlers.")
+ .isGreaterThanOrEqualTo(expectS3Handler ? 1 : 1);
+ }
+
+ } finally {
+ installer.onStop(new ShutdownEvent());
+ }
+ var onStopHandlers = installedHandlers();
+
+ // always expect that the S3 handler is removed.
+ assertThat(installedHandlers())
+ .as("onStop removes S3 handler")
+ .doesNotContain(installer.handler());
+
+ if (expectOtherHandlers) {
+ assertThat(onStopHandlers.size())
+ .as("onStart left other handler in place")
+ .isGreaterThanOrEqualTo(1);
+ } else {
+ assertThat(onStopHandlers.size()).as("onStart removed all handlers.").isEqualTo(0);
+ }
+ }
+
+ private static List installedHandlers() {
+ return Arrays.stream(Logger.getLogger(BILLING_LOGGER_NAME).getHandlers()).toList();
+ }
+
+ private static BillingS3HandlerInstaller installer(
+ Map configOverride, MeterRegistry meterRegistry) {
+
+ var builder = new SmallRyeConfigBuilder().withMapping(BillingS3ExportConfig.class);
+ if (configOverride != null) {
+ configOverride.forEach(builder::withDefaultValue);
+ }
+ var config = builder.build().getConfigMapping(BillingS3ExportConfig.class);
+
+ meterRegistry = meterRegistry == null ? new SimpleMeterRegistry() : meterRegistry;
+
+ return new BillingS3HandlerInstaller(config, meterRegistry);
+ }
+}
diff --git a/src/test/java/io/stargate/sgv2/jsonapi/service/provider/BillingTest.java b/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingTest.java
similarity index 94%
rename from src/test/java/io/stargate/sgv2/jsonapi/service/provider/BillingTest.java
rename to src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingTest.java
index 399155d241..e5b589cae9 100644
--- a/src/test/java/io/stargate/sgv2/jsonapi/service/provider/BillingTest.java
+++ b/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingTest.java
@@ -1,4 +1,4 @@
-package io.stargate.sgv2.jsonapi.service.provider;
+package io.stargate.sgv2.jsonapi.service.billing;
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatCode;
@@ -12,6 +12,10 @@
import io.stargate.sgv2.jsonapi.config.feature.ApiFeature;
import io.stargate.sgv2.jsonapi.config.feature.ApiFeatures;
import io.stargate.sgv2.jsonapi.config.feature.FeaturesConfig;
+import io.stargate.sgv2.jsonapi.service.provider.ModelInputType;
+import io.stargate.sgv2.jsonapi.service.provider.ModelProvider;
+import io.stargate.sgv2.jsonapi.service.provider.ModelType;
+import io.stargate.sgv2.jsonapi.service.provider.ModelUsage;
import io.vertx.core.MultiMap;
import java.util.List;
import java.util.Map;
diff --git a/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingTestBase.java b/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingTestBase.java
new file mode 100644
index 0000000000..a506a83db8
--- /dev/null
+++ b/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingTestBase.java
@@ -0,0 +1,484 @@
+package io.stargate.sgv2.jsonapi.service.billing;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.Mockito.*;
+
+import io.stargate.sgv2.jsonapi.metrics.BatchedLogBufferMetrics;
+import io.stargate.sgv2.jsonapi.util.MockClock;
+import java.time.Duration;
+import java.util.List;
+import java.util.concurrent.*;
+import java.util.concurrent.locks.LockSupport;
+import java.util.logging.Level;
+import java.util.logging.LogRecord;
+import java.util.stream.IntStream;
+import java.util.stream.Stream;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/** Common code for tests around billing events being uploaded */
+public abstract class BillingTestBase {
+
+ private static final Logger LOGGER = LoggerFactory.getLogger(BillingTestBase.class);
+
+ // ======================================================================
+ // BUFFER config and values for the tests to use
+ // ======================================================================
+
+ // want the line bytes when lines go into the buffer to be 25
+ // template below is 21 bytes
+ // 3 chars for the index get added in createFixture()
+ // 1 char added in the buffer calc's for the `\n` to write out
+ protected static final int MESSAGE_LENGTH_IN_BUFFER = 25;
+ protected static final String TEMPLATE_25_CHARS = "Total of 25 chars ";
+
+ protected static final int MAX_BATCH_SIZE = 100;
+ // The number of messages we can fit inside the max bytes setting
+ protected static final int MAX_BATCH_BYTES_NUM_MESSAGES = 20;
+ protected static final int MAX_BATCH_BYTES =
+ MESSAGE_LENGTH_IN_BUFFER * MAX_BATCH_BYTES_NUM_MESSAGES;
+
+ // How many full batches, tracked by max size, we want to fit in the buffer
+ protected static final int BATCHES_BY_SIZE_PER_CAPACITY = 3;
+ protected static final int BUFFER_CAPACITY = MAX_BATCH_SIZE * BATCHES_BY_SIZE_PER_CAPACITY;
+
+ // number of log records we create for each feature / test
+ protected static final int NUM_RECORDS = BUFFER_CAPACITY * 3;
+ // when using mock clock, we set the instant for each log record to be 1 "second"
+ // after the last, so we will create log records with up to
+ // NUM_RECORDS of seconds past when the clock was started
+ // used when testing the max age features
+ protected static final Duration MAX_AGE = Duration.ofSeconds(NUM_RECORDS);
+ protected static final Level LOG_LEVEL = Level.INFO;
+
+ // ======================================================================
+ // LOG HANDLER config and values for the tests to use
+ // ======================================================================
+
+ // sleep between checking the buffer, long we will normally use flush() to wake
+ // for tests
+ protected static final Duration UPLOAD_SLEEP_DURATION = Duration.ofSeconds(60);
+ protected static final Duration UPLOAD_SLEEP_DURATION_SHORT = Duration.ofMillis(100);
+
+ // upload must complete in this time
+ protected static final Duration UPLOADER_SAFETY_DEADLINE = Duration.ofSeconds(30);
+ // close() will wait this long for startUploading() to finish
+ protected static final Duration UPLOAD_SHUTDOWN_DEADLINE = Duration.ofSeconds(30);
+ protected static final Duration UPLOAD_SHUTDOWN_DEADLINE_SHORT = Duration.ofSeconds(1);
+
+ protected static void threadSleep(long millis) {
+ LockSupport.parkNanos(Duration.ofMillis(millis).toNanos());
+ }
+
+ protected static void waitOnLatch(CountDownLatch latch) {
+ try {
+ latch.await();
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new RuntimeException("interrupted waiting on latch", e);
+ }
+ }
+
+ // *********************************************************
+ // Fixture - used to group the test config and test data we need
+ // *********************************************************
+
+ /** Default fixture for buffer tests with config from the top of class */
+ Fixture defaultBufferFixture() {
+ return defaultBufferFixture(false);
+ }
+
+ Fixture defaultBufferFixture(boolean mockBufferClock) {
+ return Fixture.createFixture(
+ MAX_BATCH_SIZE,
+ MAX_BATCH_BYTES,
+ MAX_AGE,
+ BUFFER_CAPACITY,
+ NUM_RECORDS,
+ LOG_LEVEL,
+ TEMPLATE_25_CHARS,
+ mockBufferClock,
+ false,
+ false,
+ false);
+ }
+
+ Fixture defaultLogHandlerFixture() {
+ return defaultLogHandlerFixture(true, false, false);
+ }
+
+ Fixture defaultLogHandlerFixture(
+ boolean mockBuffer, boolean shortUploadSleep, boolean shortShutdownDuration) {
+ return Fixture.createFixture(
+ MAX_BATCH_SIZE,
+ MAX_BATCH_BYTES,
+ MAX_AGE,
+ BUFFER_CAPACITY,
+ NUM_RECORDS,
+ LOG_LEVEL,
+ TEMPLATE_25_CHARS,
+ false,
+ mockBuffer,
+ shortUploadSleep,
+ shortShutdownDuration);
+ }
+
+ /**
+ * Tracks the config of the buffer, the buffer, the data we can use for each test to add to
+ * buffer, etc.
+ *
+ * See {@link #defaultBufferFixture(boolean)}
+ */
+ record Fixture(
+ int maxBatchSize,
+ long maxBytes,
+ Duration maxAge,
+ int queueCapacity,
+ List logRecords,
+ BatchedLogBuffer buffer,
+ BatchedLogBufferMetrics metrics,
+ BillingUploadingLogHandler logHandler,
+ AsyncBatchedLogUploader uploader,
+ MockClock clock) {
+
+ /** Create fixture, creates LogRecords that can be used to add to the buffer */
+ static Fixture createFixture(
+ int maxBatchSize,
+ long maxBytes,
+ Duration maxAge,
+ int queueCapacity,
+ int numLogRecords,
+ Level logLevel,
+ String logRecordTemplate,
+ boolean mockBufferClock,
+ boolean mockBuffer,
+ boolean shortUploadSleepDuration,
+ boolean shortUploadShutdownDeadline) {
+
+ if (mockBuffer && mockBufferClock) {
+ throw new IllegalArgumentException("cannot mock the buffer and the buffer clock");
+ }
+ // Make sure to initialize the mock clock before creating the log messages
+ // so they are always after the start of the clock.
+ var mockClock = mockBufferClock ? new MockClock() : null;
+
+ // fork the clock, we are going to use clockForRecords when creating the records
+ // and will advance it 1 second for each record, the original mockClock is for
+ // the buffer to use, so we let the test advance that
+ var clockForRecords = mockClock == null ? null : new MockClock(mockClock);
+
+ var logRecords =
+ IntStream.range(0, numLogRecords)
+ .mapToObj(i -> logRecordTemplate + String.format("%03d", i))
+ .map(
+ s -> {
+ var record = new LogRecord(logLevel, s);
+ if (clockForRecords != null) {
+ record.setInstant(clockForRecords.instant());
+ clockForRecords.nextSecond();
+ }
+ return record;
+ })
+ .toList();
+
+ var metrics = mock(BatchedLogBufferMetrics.class);
+
+ BatchedLogBuffer buffer;
+ if (mockBuffer) {
+ buffer = mock(BatchedLogBuffer.class);
+ // setup for an empty buffer when calling
+ when(buffer.offer(any())).thenReturn(true);
+ when(buffer.isEmpty()).thenReturn(true);
+ when(buffer.size()).thenReturn(0);
+ when(buffer.nextBatch(anyBoolean())).thenReturn(null);
+ } else {
+ buffer =
+ new BatchedLogBuffer(
+ maxBatchSize,
+ maxBytes,
+ maxAge,
+ queueCapacity,
+ metrics,
+ mockBufferClock ? mockClock : BatchedLogBuffer.DEFAULT_CLOCK);
+ }
+ var uploader = mock(AsyncBatchedLogUploader.class);
+ var logHandler =
+ new BillingUploadingLogHandler(
+ buffer,
+ uploader,
+ shortUploadSleepDuration ? UPLOAD_SLEEP_DURATION_SHORT : UPLOAD_SLEEP_DURATION,
+ UPLOADER_SAFETY_DEADLINE,
+ shortUploadShutdownDeadline
+ ? UPLOAD_SHUTDOWN_DEADLINE_SHORT
+ : UPLOAD_SHUTDOWN_DEADLINE);
+
+ return new Fixture(
+ maxBatchSize,
+ maxBytes,
+ maxAge,
+ queueCapacity,
+ logRecords,
+ buffer,
+ metrics,
+ logHandler,
+ uploader,
+ mockClock);
+ }
+
+ /** Assert the buffer is full, and so offer() fails */
+ void assertBufferFull(String desc, int index) {
+
+ // although the next log record is wafer-thin, it is too much for Mr Creosote
+ assertThat(buffer().offer(logRecords.get(index))).as(desc + " - fail at capacity").isFalse();
+
+ // Running again to confirm it is still full
+ assertThat(buffer().offer(logRecords.get(index)))
+ .as(desc + " - second - fail at capacity")
+ .isFalse();
+ }
+
+ /**
+ * Offer the log records selected by slice to the buffer, all should work, assert the buffer has
+ * the items the slice selected
+ */
+ void assertOffer(String desc, BatchedLogBufferTest.Slice slice) {
+
+ var snapshot = BufferSnapshot.create(this);
+
+ for (var record : slice.stream(logRecords).toList()) {
+ assertThat(buffer.offer(record)).as(desc + " - assertOffer() - offering").isTrue();
+ }
+
+ snapshot.assertAll(desc, slice, true);
+ }
+
+ /**
+ * Get a batch from the buffer, assert we got a batch that is legal, and assert the buffer has
+ * changed by the amount of the batch
+ */
+ BatchedLogBuffer.Batch assertNextBatch(String desc, boolean drainFully) {
+
+ var snapshot = BufferSnapshot.create(this);
+ var batch = buffer.nextBatch(drainFully);
+
+ // assert the batch is what we expected.
+ assertThat(batch).as(desc + " - assertNextBatch() - batch is not null").isNotNull();
+
+ assertThat(batch.size())
+ .as(desc + " - assertNextBatch() - batch size <= MAX_BATCH_SIZE")
+ .isLessThanOrEqualTo(maxBatchSize);
+ // note: it is legal to have a batch bigger than the maxBytes, specialised tests for that
+ // shoudl only happen when there is a single log record bigger than maxBytes
+ assertThat(batch.bytes())
+ .as(desc + " - assertNextBatch() - batch bytes <= MAX_BATCH_BYTES")
+ .isLessThanOrEqualTo(maxBytes);
+
+ // assert the buffer updated bookkeeping as we expect
+ snapshot.assertAll(desc, batch);
+ return batch;
+ }
+
+ // just redeclare with no exception to make it easier
+ interface NoExceptionCloseable extends AutoCloseable {
+ @Override
+ void close();
+ }
+
+ /**
+ * Start the logHandler uploading on a daemon thread, and returns a closeable for killing the
+ * thread.
+ *
+ * @return
+ */
+ NoExceptionCloseable startHandlerUploading(String desc) {
+
+ var executor =
+ Executors.newSingleThreadExecutor(Thread.ofPlatform().daemon().name(desc).factory());
+
+ var uploadingFuture =
+ executor.submit(
+ () -> {
+ LOGGER.info("startHandlerUploading() - starting logHandler. desc:{}", desc);
+ logHandler.startUploading();
+ });
+
+ return () -> {
+ LOGGER.info("startHandlerUploading() - stopping logHandler. desc:{}", desc);
+ try {
+ // this is waiting for the uploading thread to return
+ uploadingFuture.get(10, TimeUnit.SECONDS);
+ executor.awaitTermination(1000, TimeUnit.MILLISECONDS);
+
+ LOGGER.info("startHandlerUploading() - stopped logHandler. desc:{}", desc);
+
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ LOGGER.error("startHandlerUploading() - error 1", e);
+ throw new RuntimeException(e);
+ } catch (ExecutionException e) {
+ LOGGER.error("startHandlerUploading() - error 2", e);
+
+ throw e.getCause() instanceof RuntimeException re
+ ? re
+ : new RuntimeException(e.getCause());
+ } catch (TimeoutException e) {
+ LOGGER.error("startHandlerUploading() - error 3", e);
+
+ throw new RuntimeException(e);
+ } finally {
+ executor.shutdownNow();
+ }
+ };
+ }
+ }
+
+ /**
+ * A slice of a list, `from` is inclusive, `to` is exclusive
+ *
+ * ...
+ */
+ record Slice(int from, int to) {
+
+ public Stream stream(List list) {
+ return list.stream().skip(from).limit(to - from);
+ }
+
+ public int size() {
+ return to - from;
+ }
+
+ public static Slice to(int to) {
+ return new Slice(0, to);
+ }
+
+ public static Slice from(int from) {
+ return new Slice(from, Integer.MAX_VALUE);
+ }
+
+ public static Slice slice(int from, int to) {
+ return new Slice(from, to);
+ }
+ }
+
+ /**
+ * Snapshot of the metadata (size etc) for the buffer, that can be used to compare how the buffer
+ * metadata has changed
+ *
+ * ...
+ */
+ record BufferSnapshot(
+ boolean isEmpty, int size, long queuedBytes, int remainingCapacity, Fixture fixture) {
+
+ static BufferSnapshot create(Fixture fixture) {
+ // reset the counters for calls to metrics
+ clearInvocations(fixture.metrics);
+ return new BufferSnapshot(
+ fixture.buffer().isEmpty(),
+ fixture.buffer().size(),
+ fixture.buffer().queuedBytes(),
+ fixture.buffer().remainingCapacity(),
+ fixture);
+ }
+
+ /**
+ * Assert that the current metadata values for the buffer are the values in the snapshot PLUS
+ * the log records that were added by the Slice.
+ */
+ void assertAll(String desc, Slice slice, boolean inOrder) {
+ assertBufferMetadata(desc, slice);
+ assertBufferItems(desc, slice, inOrder);
+ }
+
+ /**
+ * Assert that the current metadata values for the buffer are the values in the snapshot MINUS
+ * the buffer entries that were removed in the batch
+ */
+ void assertAll(String desc, BatchedLogBuffer.Batch batch) {
+ assertBufferMetadata(desc, batch);
+ assertBufferItems(desc, batch);
+ }
+
+ /** current buffer metadata = snapshot + slice */
+ void assertBufferMetadata(String desc, Slice slice) {
+
+ if (slice.size() == 0) {
+ assertThat(fixture.buffer().isEmpty())
+ .as(desc + " - isEmpty no change after empty slice")
+ .isEqualTo(isEmpty());
+ } else {
+ assertThat(fixture.buffer().isEmpty())
+ .as(desc + " - isEmpty false after non empty slice")
+ .isEqualTo(false);
+ }
+
+ assertThat(fixture.buffer().size())
+ .as(desc + " - post buffer size increased by slice")
+ .isEqualTo(size() + slice.size());
+
+ verify(
+ fixture.metrics,
+ times(slice.size()).description(desc + "metrics called for every offer"))
+ .offered();
+
+ long addedBytes = 0;
+ for (var record : slice.stream(fixture.logRecords).toList()) {
+ addedBytes += BatchedLogBuffer.Entry.lineBytes(record.getMessage());
+ }
+
+ assertThat(fixture.buffer().queuedBytes())
+ .as(desc + " - post buffer bytes increased by slice")
+ .isEqualTo(queuedBytes + addedBytes);
+ }
+
+ /** current buffer metadata = snapshot - batch */
+ void assertBufferMetadata(String desc, BatchedLogBuffer.Batch batch) {
+
+ assertThat(fixture.buffer().size())
+ .as(desc + " - buffer size decreased by batch size")
+ .isEqualTo(size() - batch.size());
+
+ assertThat(fixture.buffer().queuedBytes())
+ .as(desc + " - buffer bytes size decreased by batch bytes")
+ .isEqualTo(queuedBytes - batch.bytes());
+ }
+
+ /**
+ * current buffer items contain items from slice inOrder - if we expect items in buffer to match
+ * order of the fixture
+ */
+ void assertBufferItems(String desc, Slice slice, boolean inOrder) {
+
+ var bufferItems = fixture.buffer().peekBuffer();
+
+ int i = slice.from() > bufferItems.size() ? 0 : slice.from();
+ for (var record : slice.stream(fixture.logRecords).toList()) {
+
+ if (inOrder) {
+ assertThat(record.getMessage())
+ .as(desc + " - buffer items at position match exactly pos: " + i)
+ .isEqualTo(bufferItems.get(i++).line());
+ } else {
+
+ var entry = new BatchedLogBuffer.Entry(record.getInstant(), record.getMessage());
+ assertThat(bufferItems)
+ .as(desc + " - buffer items contains entry: " + entry)
+ .contains(entry);
+ }
+ }
+ }
+
+ /** current buffer items contain NONE of items in batch */
+ void assertBufferItems(String desc, BatchedLogBuffer.Batch batch) {
+
+ var peekedBuffer = fixture.buffer().peekBuffer();
+
+ for (var batchString : batch.lines()) {
+
+ var found = peekedBuffer.stream().anyMatch(entry -> entry.line().equals(batchString));
+ assertThat(found)
+ .as(desc + " - line from batch no longer in buffer: " + batchString)
+ .isFalse();
+ }
+ }
+ }
+}
diff --git a/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingUploadingLogHandlerTest.java b/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingUploadingLogHandlerTest.java
new file mode 100644
index 0000000000..7113c62967
--- /dev/null
+++ b/src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingUploadingLogHandlerTest.java
@@ -0,0 +1,309 @@
+package io.stargate.sgv2.jsonapi.service.billing;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+import static org.mockito.Mockito.*;
+
+import io.smallrye.mutiny.Uni;
+import java.util.List;
+import java.util.logging.LogRecord;
+import java.util.stream.IntStream;
+import org.junit.jupiter.api.Test;
+import org.mockito.ArgumentCaptor;
+import org.mockito.verification.VerificationMode;
+
+/** */
+public class BillingUploadingLogHandlerTest extends BillingTestBase {
+
+ // *********************************************************
+ // Handler interface - Producer side of the handler
+ // *********************************************************
+
+ /** Null record silently dropped by handler */
+ @Test
+ public void publishSilentDropNulls() {
+
+ var fixture = defaultLogHandlerFixture();
+
+ fixture.logHandler().publish(null);
+ fixture.logHandler().publish(null);
+ fixture.logHandler().publish(null);
+
+ verify(
+ fixture.buffer(),
+ times(0).description("publishSilentDropNulls() - no calls to buffer.offer()"))
+ .offer(any());
+ }
+
+ /** Non null record silently dropped by handler when closed */
+ @Test
+ public void publishSilentDropWhenClosed() {
+
+ var fixture = defaultLogHandlerFixture();
+ fixture.logHandler().close();
+
+ fixture.logHandler().publish(fixture.logRecords().getFirst());
+ fixture.logHandler().publish(fixture.logRecords().getFirst());
+ fixture.logHandler().publish(fixture.logRecords().getFirst());
+
+ verify(
+ fixture.buffer(),
+ times(0).description("publishSilentDropWhenClosed() - no calls to buffer.offer()"))
+ .offer(any());
+ }
+
+ /** Records published when buffer is full are dropped, no error */
+ @Test
+ public void publishSilentWhenBufferFull() {
+
+ var fixture = defaultLogHandlerFixture();
+
+ // return false, buffer full , go away
+ when(fixture.buffer().offer(any())).thenReturn(false);
+
+ fixture.logHandler().publish(fixture.logRecords().getFirst());
+ fixture.logHandler().publish(fixture.logRecords().getFirst());
+ fixture.logHandler().publish(fixture.logRecords().getFirst());
+
+ verify(
+ fixture.buffer(),
+ times(3).description("publishSilentWhenBufferFull() - called offer for each record"))
+ .offer(fixture.logRecords().getFirst());
+ }
+
+ /** Passing record to handler, is then passed to the buffer. */
+ @Test
+ public void publishOfferSucceed() {
+
+ var fixture = defaultLogHandlerFixture();
+ var slice = Slice.to(MAX_BATCH_BYTES_NUM_MESSAGES);
+ var expectedLogRecords = slice.stream(fixture.logRecords()).toList();
+
+ var argCaptor = ArgumentCaptor.forClass(LogRecord.class);
+
+ for (var record : expectedLogRecords) {
+ fixture.logHandler().publish(record);
+ }
+
+ verify(
+ fixture.buffer(),
+ times(expectedLogRecords.size())
+ .description("publishOfferSucceed() - buffer called for each log record"))
+ .offer(argCaptor.capture());
+ var actualLogRecords = argCaptor.getAllValues();
+
+ assertThat(actualLogRecords)
+ .as("publishOfferSucceed() - all and only expected records passed to the buffer")
+ .containsExactlyElementsOf(expectedLogRecords);
+ }
+
+ /** Calling flush on handler that is NOT uploading does nothing */
+ @Test
+ public void flushNullOpIfNotStarted() {
+
+ var fixture = defaultLogHandlerFixture();
+
+ fixture.logHandler().flush();
+ // there is no uploading, so should not ask for next batch
+ verify(fixture.buffer(), never()).nextBatch(anyBoolean());
+ }
+
+ /**
+ * Verify the number of times a function was called, but with a timeout to wait. e.g. when waiting
+ * for the uploading thread to wakeup
+ */
+ private VerificationMode timeoutTimes(String desc, int times) {
+ return timeout(2000).times(times).description(desc);
+ }
+
+ /** Calling flush on handler that is uploading causes handler to check buffer. */
+ @Test
+ public void flushChecksForBatch() {
+
+ var fixture1 = defaultLogHandlerFixture();
+ try (var handlerThread =
+ fixture1.startHandlerUploading("flushChecksForBatch() - close not called")) {
+ fixture1.logHandler().flush();
+ // close has not been called, so it should not drain
+ verify(fixture1.buffer(), timeoutTimes("nextBatch() called once with drain false", 1))
+ .nextBatch(false);
+
+ // this is a bit stupid, calling close to close the upload thread
+ fixture1.logHandler().close();
+ }
+
+ var fixture2 = defaultLogHandlerFixture();
+ try (var handlerThread =
+ fixture2.startHandlerUploading("flushChecksForBatch() - close is called")) {
+ fixture2.logHandler().unsafeClose();
+ fixture2.logHandler().flush();
+ // close has been called, so it should drain buffer
+ verify(fixture2.buffer(), timeoutTimes("nextBatch() called once with drain true", 1))
+ .nextBatch(true);
+ }
+ }
+
+ @Test
+ public void closeWithoutUploadThreadReturns() {
+
+ var fixture1 = defaultLogHandlerFixture();
+ // NOT STARTING upload
+ fixture1.logHandler().close();
+ // upload not running, should not try to get a batch
+ verify(fixture1.buffer(), timeoutTimes("nextBatch() never called", 0)).nextBatch(anyBoolean());
+ // should have closed the uploader
+ verify(fixture1.uploader(), timeoutTimes("uploader.close() called", 1)).close();
+ }
+
+ /**
+ * Call close, but the upload thread has not released the upload permit, so close cannot detect
+ * upload has finished.
+ */
+ @Test
+ public void closeReturnsWhenUploadUnstopped() {
+
+ var fixture1 = defaultLogHandlerFixture(true, false, true);
+ // NOT STARTING upload, but acquire the permit it would take
+ fixture1.logHandler().unsafeAcquireUploadPermit();
+ // close() will not return until it times out waiting for the upload thread to finish
+ fixture1.logHandler().close();
+ }
+
+ /**
+ * Calling close() when the handler is running should cause the buffer to be called to drain it.
+ */
+ @Test
+ public void closeCausesBufferDrain() {
+
+ var fixture1 = defaultLogHandlerFixture();
+ try (var handlerThread = fixture1.startHandlerUploading("closeCausesBufferDrain()")) {
+ threadSleep(100); // give the uploader time to get into the wait on wakeup
+
+ fixture1.logHandler().close();
+ // close has been called, so it should drain
+ verify(fixture1.buffer(), timeoutTimes("nextBatch() called with drain=true", 1))
+ .nextBatch(true);
+ verify(fixture1.uploader(), timeoutTimes("uploader.close() called", 1)).close();
+ }
+ // the auto closable will wait for the upload thread to naturally exit
+ }
+
+ // *********************************************************
+ // startUploading - Consumer side of the handler
+ // *********************************************************
+
+ /**
+ * Calling startUpLoad twice on different threads, fails because there can be only one active
+ * thread running the function
+ *
+ *
Cannot call on same thread as it will be parked running the upload
+ */
+ @Test
+ public void startUploadingCalledTwiceFails() {
+
+ var fixture1 = defaultLogHandlerFixture();
+ try (var handlerThread1 =
+ fixture1.startHandlerUploading("startUploadingCalledTwiceFails() - 1st")) {
+ // make sure the worker thread has time to start
+ threadSleep(10);
+
+ // the exception will happen when startUploading is entered, but we
+ // wont get the error until calling close() which calls Future.get()
+ var closable = fixture1.startHandlerUploading("startUploadingCalledTwiceFails() - 2nd");
+ // make sure the worker thread has time to start
+ threadSleep(10);
+
+ assertThatThrownBy(closable::close, "startUploadingCalledTwiceFails() - second call")
+ .isInstanceOf(IllegalStateException.class);
+
+ // stop the first thread that is running startUpload()
+ fixture1.logHandler().close();
+ }
+ }
+
+ /** In normal operation startUpload detects three batches and sends to uploader */
+ @Test
+ public void startUploadingSendsToUploader() {
+
+ var NUM_BATCHES = 3;
+ var fixture1 = defaultLogHandlerFixture();
+ try (var handlerThread1 =
+ fixture1.startHandlerUploading("startUploadingSendsToUploader() - upload thread")) {
+ // make sure the worker thread has time to start and get to the sleep.
+ threadSleep(100);
+
+ // ** TESTING NORMAL OPERATION
+
+ // setup buffer to return three batches we will collect in normal operations
+ var expectedNormalBatches = mockUploading(fixture1, NUM_BATCHES, false);
+ // handler should be sleeping because of long sleep, flush will wake it up.
+ fixture1.logHandler().flush();
+ // wait for it the handler to call the uploader
+ var normalCaptor = ArgumentCaptor.forClass(BatchedLogBuffer.Batch.class);
+ verify(
+ fixture1.uploader(),
+ timeoutTimes(
+ "startUploadingSendsToUploader() - normal mode", expectedNormalBatches.size()))
+ .upload(normalCaptor.capture());
+ var actualNormalBatches = normalCaptor.getAllValues();
+
+ // ** TESTING CLOSE / SHUTDOWN OPERATION
+
+ // reset counter , will already be NUM_BATCHES from the normal operation check above
+ clearInvocations(fixture1.uploader());
+ // we are using the long upload sleep, uploader should be asleep again, send batches
+ // and close to see we get correct behavior
+ var expectedShutdownBatches = mockUploading(fixture1, NUM_BATCHES, true);
+ // close to get shutdown operations
+ fixture1.logHandler().close();
+ // wait for it the handler to call the uploader
+ var shutdownCaptor = ArgumentCaptor.forClass(BatchedLogBuffer.Batch.class);
+ verify(
+ fixture1.uploader(),
+ timeoutTimes(
+ "startUploadingSendsToUploader() - shutdown mode",
+ expectedShutdownBatches.size()))
+ .upload(shutdownCaptor.capture());
+ var actualShutdownBatches = shutdownCaptor.getAllValues();
+
+ assertThat(actualNormalBatches)
+ .as("startUploadingSendsToUploader() - batches from normal operation match")
+ .containsExactlyElementsOf(expectedNormalBatches);
+
+ assertThat(actualShutdownBatches)
+ .as("startUploadingSendsToUploader() - batches from shutdown operation match")
+ .containsExactlyElementsOf(expectedShutdownBatches);
+
+ // we have closed handler, should exit block now
+ }
+ }
+
+ private List mockUploading(
+ Fixture fixture, int numBatches, boolean expectDrainFully) {
+ // by default the mock buffer will be returning null for nextBatch(), the
+ // log handler should be in a wait because we are using long upload sleep
+
+ var expectedBatches =
+ IntStream.range(0, numBatches).mapToObj(i -> mock(BatchedLogBuffer.Batch.class)).toList();
+
+ // we expect the uploader to be called with these batches, and it needs to
+ // return a Uni with an UploadResult - but we dont need to keep the UploadResult
+ // MUST set up the uploader to return a value before putting batches into the buffer
+ for (var batch : expectedBatches) {
+ var result = new AsyncBatchedLogUploader.UploadResult(batch, null);
+ when(fixture.uploader().upload(batch)).thenReturn(Uni.createFrom().item(result));
+ }
+
+ // now connect the Batch's to the buffer so it will return them. Once this is done
+ // the upload thread will pick them up once woken
+ // expectDrainFully depends on if the test is closed the handler.
+ var nextBatchStub = when(fixture.buffer().nextBatch(expectDrainFully));
+ for (var batch : expectedBatches) {
+ nextBatchStub = nextBatchStub.thenReturn(batch);
+ }
+ // now return null as sentinel for no more data
+ nextBatchStub.thenReturn(null);
+
+ return expectedBatches;
+ }
+}
diff --git a/src/test/java/io/stargate/sgv2/jsonapi/service/provider/DefaultBillingTest.java b/src/test/java/io/stargate/sgv2/jsonapi/service/billing/DefaultBillingTest.java
similarity index 97%
rename from src/test/java/io/stargate/sgv2/jsonapi/service/provider/DefaultBillingTest.java
rename to src/test/java/io/stargate/sgv2/jsonapi/service/billing/DefaultBillingTest.java
index 7657c48282..30b570ced6 100644
--- a/src/test/java/io/stargate/sgv2/jsonapi/service/provider/DefaultBillingTest.java
+++ b/src/test/java/io/stargate/sgv2/jsonapi/service/billing/DefaultBillingTest.java
@@ -1,4 +1,4 @@
-package io.stargate.sgv2.jsonapi.service.provider;
+package io.stargate.sgv2.jsonapi.service.billing;
import static java.util.logging.Logger.getLogger;
import static net.javacrumbs.jsonunit.JsonAssert.assertJsonEquals;
@@ -9,6 +9,10 @@
import io.stargate.sgv2.jsonapi.TestConstants;
import io.stargate.sgv2.jsonapi.config.BillingConfig;
+import io.stargate.sgv2.jsonapi.service.provider.ModelInputType;
+import io.stargate.sgv2.jsonapi.service.provider.ModelProvider;
+import io.stargate.sgv2.jsonapi.service.provider.ModelType;
+import io.stargate.sgv2.jsonapi.service.provider.ModelUsage;
import java.time.Instant;
import java.util.ArrayList;
import java.util.Arrays;
diff --git a/src/test/java/io/stargate/sgv2/jsonapi/service/billing/S3BatchedLogUploaderTest.java b/src/test/java/io/stargate/sgv2/jsonapi/service/billing/S3BatchedLogUploaderTest.java
new file mode 100644
index 0000000000..e0802e4921
--- /dev/null
+++ b/src/test/java/io/stargate/sgv2/jsonapi/service/billing/S3BatchedLogUploaderTest.java
@@ -0,0 +1,284 @@
+package io.stargate.sgv2.jsonapi.service.billing;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.Mockito.*;
+
+import io.smallrye.config.SmallRyeConfigBuilder;
+import io.stargate.sgv2.jsonapi.config.BillingS3ExportConfig;
+import io.stargate.sgv2.jsonapi.metrics.BatchedLogUploaderMetrics;
+import java.nio.charset.StandardCharsets;
+import java.time.Duration;
+import java.time.Instant;
+import java.time.temporal.ChronoUnit;
+import java.util.List;
+import java.util.concurrent.CompletableFuture;
+import org.junit.jupiter.api.Test;
+import org.mockito.ArgumentCaptor;
+import software.amazon.awssdk.core.async.AsyncRequestBody;
+import software.amazon.awssdk.http.SdkHttpResponse;
+import software.amazon.awssdk.services.s3.S3AsyncClient;
+import software.amazon.awssdk.services.s3.model.PutObjectRequest;
+import software.amazon.awssdk.services.s3.model.PutObjectResponse;
+
+/**
+ * The S3 client is mocked; retries and per-call timeouts live in the client configuration, so
+ * exactly one {@code putObject} per upload is expected here. Real I/O is covered by {@code
+ * BillingS3ExportIntegrationTest}.
+ */
+class S3BatchedLogUploaderTest {
+
+ //
+ // private static final Pattern KEY_PATTERN =
+ // Pattern.compile("data-api/2026/05/20/14/23/[0-9a-f-]{36}\\.jsonl");
+ 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());
+
+ // private static S3BatchedLogUploader uploader(S3AsyncClient client) {
+ // return new S3BatchedLogUploader(client, "my-bucket");
+ // }
+ //
+ // private static CompletableFuture ok() {
+ // return CompletableFuture.completedFuture(PutObjectResponse.builder().build());
+ // }
+
+ @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);
+ }
+
+ // ============================================================
+ // lifecycle + config validation
+ // ============================================================
+
+ // @Test
+ // void closeClosesTheClient() {
+ // S3AsyncClient client = mock(S3AsyncClient.class);
+ // uploader(client).close();
+ // verify(client).close();
+ // }
+ //
+ // @Test
+ // void createRejectsMissingRegionOrBucket() {
+ // assertThatThrownBy(() -> S3BatchedLogUploader.create(" ", "bucket", Optional.empty()))
+ // .isInstanceOf(IllegalArgumentException.class)
+ // .hasMessageContaining("bucket-region");
+ // assertThatThrownBy(() -> S3BatchedLogUploader.create("us-east-1", null, Optional.empty()))
+ // .isInstanceOf(IllegalArgumentException.class)
+ // .hasMessageContaining("billing.s3.bucket");
+ // }
+
+ // ============================================================
+ // Scaffold
+ // ============================================================
+
+ private record Fixture(
+ S3BatchedLogUploader uploader,
+ S3AsyncClient s3Client,
+ PutObjectResponse putResponse,
+ BatchedLogUploaderMetrics metrics,
+ BillingS3ExportConfig config) {
+
+ static Fixture create() {
+ return create(false);
+ }
+
+ static Fixture create(boolean useFactory) {
+
+ var config =
+ new SmallRyeConfigBuilder()
+ .withMapping(BillingS3ExportConfig.class)
+ .build()
+ .getConfigMapping(BillingS3ExportConfig.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..9ef3c89e76
--- /dev/null
+++ b/src/test/java/io/stargate/sgv2/jsonapi/testresource/S3MockTestResource.java
@@ -0,0 +1,89 @@
+package io.stargate.sgv2.jsonapi.testresource;
+
+import com.adobe.testing.s3mock.testcontainers.S3MockContainer;
+import io.quarkus.test.common.QuarkusTestResourceLifecycleManager;
+import java.util.HashMap;
+import java.util.Map;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Starts an S3Mock container and enables the billing S3 export against it, with small batch
+ * thresholds so tests see objects quickly. Used by {@code BillingS3ExportIntegrationTest} alongside
+ * {@link DseTestResource}.
+ *
+ * The returned properties reach the application under test (a separate process for
+ * {@code @QuarkusIntegrationTest}); mirroring them as system properties follows the {@link
+ * StargateTestResource} pattern so the whole test environment sees the same values.
+ */
+public class S3MockTestResource implements QuarkusTestResourceLifecycleManager {
+
+ private static final Logger LOG = LoggerFactory.getLogger(S3MockTestResource.class);
+
+ /** Container tag; keep in sync with the {@code s3mock-testcontainers} version in pom.xml. */
+ private static final String S3MOCK_VERSION = "5.1.0";
+
+ public static final String BUCKET = "billing-events-it";
+ public static final String BUCKET_REGION = "us-east-1";
+ public static final String ACCESS_KEY = "s3mock-test";
+ public static final String SECRET_KEY = "s3mock-test";
+
+ private static volatile String httpEndpoint;
+
+ private 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, like an S3 outage. One-way for the whole test class
+ * (a restart would map a new port, unreachable through the app's fixed endpoint-override), so
+ * only the last test may call this.
+ */
+ public static void stopContainer() {
+ if (container == null) {
+ throw new IllegalStateException("S3MockTestResource has not been started");
+ }
+ container.stop();
+ }
+
+ @Override
+ public Map start() {
+ container = new S3MockContainer(S3MOCK_VERSION).withInitialBuckets(BUCKET);
+ container.start();
+ httpEndpoint = container.getHttpEndpoint();
+
+ Map props = new HashMap<>();
+ props.put("stargate.jsonapi.billing.s3.enabled", "true");
+ props.put("stargate.jsonapi.billing.s3.bucket", BUCKET);
+ props.put("stargate.jsonapi.billing.s3.bucket-region", BUCKET_REGION);
+ props.put("stargate.jsonapi.billing.s3.endpoint-override", httpEndpoint);
+ // Small thresholds so the export flushes promptly: count seal at 5, age sweep every 2s.
+ props.put("stargate.jsonapi.billing.s3.max-events", "5");
+ props.put("stargate.jsonapi.billing.s3.max-age", "PT2S");
+ props.put("stargate.jsonapi.billing.s3.shutdown-timeout", "PT5S");
+ // The producer side (DefaultBilling) is feature-flagged off by default.
+ props.put("stargate.feature.flags.billing-events-logging", "true");
+ // The uploader resolves credentials from the SDK default chain, whose first stop is the
+ // system-property provider. S3Mock accepts any signed request.
+ props.put("aws.accessKeyId", ACCESS_KEY);
+ props.put("aws.secretAccessKey", SECRET_KEY);
+
+ props.forEach(System::setProperty);
+ LOG.info("S3Mock started for billing export IT: endpoint={}, bucket={}", httpEndpoint, BUCKET);
+ return props;
+ }
+
+ @Override
+ public void stop() {
+ if (container != null) {
+ container.stop();
+ }
+ }
+}
diff --git a/src/test/java/io/stargate/sgv2/jsonapi/util/MockClock.java b/src/test/java/io/stargate/sgv2/jsonapi/util/MockClock.java
new file mode 100644
index 0000000000..a08e4baa0b
--- /dev/null
+++ b/src/test/java/io/stargate/sgv2/jsonapi/util/MockClock.java
@@ -0,0 +1,62 @@
+package io.stargate.sgv2.jsonapi.util;
+
+import java.time.*;
+import java.util.concurrent.atomic.AtomicReference;
+
+/** Implementation of the Java Clock that can be used to control time for time dependant tests. */
+public class MockClock extends Clock {
+ private final Instant startedAt;
+ private final AtomicReference now;
+ private final ZoneId zone;
+
+ public MockClock() {
+ this(Instant.now(), ZoneId.systemDefault());
+ }
+
+ public MockClock(MockClock other) {
+ this(other.startedAt, other.zone);
+ }
+
+ private MockClock(Instant now, ZoneId zone) {
+ this.now = new AtomicReference<>(now);
+ this.startedAt = now;
+ this.zone = zone;
+ }
+
+ public Instant startedAt() {
+ return startedAt;
+ }
+
+ public MockClock nextSecond() {
+ return addSeconds(1);
+ }
+
+ public MockClock addSeconds(int seconds) {
+ return advance(Duration.ofSeconds(seconds));
+ }
+
+ public MockClock advance(Duration amount) {
+ now.updateAndGet(current -> current.plus(amount));
+ return this;
+ }
+
+ public MockClock setInstant(Instant instant) {
+ now.set(instant);
+ return this;
+ }
+
+ @Override
+ public ZoneId getZone() {
+ return zone;
+ }
+
+ @Override
+ public Clock withZone(ZoneId zone) {
+ return new MockClock(now.get(), zone);
+ }
+
+ @Override
+ public Instant instant() {
+ return now.get();
+ }
+}