There are two users of this class: + * + *
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
+ * 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 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..50f2ea681c
--- /dev/null
+++ b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingS3HandlerInstaller.java
@@ -0,0 +1,100 @@
+package io.stargate.sgv2.jsonapi.service.billing;
+
+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;
+
+/**
+ * Attaches a {@link BillingUploadingLogHandler} to the {@code billing.events} logger at startup
+ * (when {@link BillingS3ExportConfig#enabled()} is {@code true}) and removes + closes it on
+ * shutdown for a graceful drain.
+ */
+@ApplicationScoped
+public class BillingS3HandlerInstaller {
+
+ private static final org.slf4j.Logger LOGGER =
+ LoggerFactory.getLogger(BillingS3HandlerInstaller.class);
+
+ private static final String METRICS_PREFIX = "billing";
+ private static final String BILLING_LOGGER_NAME = "billing.events";
+
+ private final BillingS3ExportConfig config;
+ private final MeterRegistry meterRegistry;
+
+ private volatile BillingUploadingLogHandler 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");
+
+ // Fail-loud: invalid billing S3 config throws here, aborting application startup.
+ var uploader =
+ S3BatchedLogUploader.create(
+ config.region(),
+ config.bucket(),
+ config.endpointOverride().orElse(null),
+ new BatchedLogUploaderMetrics(meterRegistry, METRICS_PREFIX));
+ LOGGER.info("onStart() - using uploader: {}", uploader);
+
+ var buffer =
+ new BatchedLogBuffer(
+ config.maxBatchSize(),
+ config.maxBatchBytes(),
+ config.maxBatchAge(),
+ config.queueCapacity(),
+ new BatchedLogBufferMetrics(meterRegistry, METRICS_PREFIX));
+ LOGGER.info("onStart() - using log buffer: {}", buffer);
+
+ this.handler =
+ new BillingUploadingLogHandler(
+ buffer,
+ uploader,
+ config.uploadSleepDuration(),
+ config.uploaderSafetyDeadline(),
+ config.uploadShutdownDeadline());
+ LOGGER.info("onStart() - using uploader: {}", uploader);
+
+ Logger.getLogger(BILLING_LOGGER_NAME).addHandler(this.handler);
+ LOGGER.info(
+ "onStart() - attached log handler to logger. BILLING_LOGGER_NAME: {}", BILLING_LOGGER_NAME);
+
+ Infrastructure.getDefaultWorkerPool().execute(this.handler::startUploading);
+ }
+
+ void onStop(@Observes ShutdownEvent event) {
+
+ if (this.handler == null) {
+ return;
+ }
+
+ // TODO: XXX WHY DO THIS ?
+ Logger.getLogger(BILLING_LOGGER_NAME).removeHandler(this.handler);
+
+ // close() isn't expected to throw, but if it does (e.g. client.close() failing), letting it
+ // propagate would disrupt other components' cleanup in Quarkus's shutdown sequence.
+ try {
+ this.handler.close();
+ } catch (Exception e) {
+ LOGGER.warn("Error during billing S3 export handler shutdown", 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..9934dc9468
--- /dev/null
+++ b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingUploadingLogHandler.java
@@ -0,0 +1,345 @@
+package io.stargate.sgv2.jsonapi.service.billing;
+
+import com.google.common.annotations.VisibleForTesting;
+import io.smallrye.mutiny.Uni;
+import java.time.Duration;
+import java.util.Objects;
+import java.util.concurrent.Semaphore;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.logging.Handler;
+import java.util.logging.LogRecord;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A Logging handler designed to be used with the Billing system. It accepts billing event log
+ * messages, batches them, and then sends to S3.
+ *
+ * See {@link BillingS3HandlerInstaller} for setup.
+ *
+ * // AI SLOP BELOW JUL handler that turns {@code billing.events} log lines into batched S3
+ * objects.
+ *
+ * Division of labor: {@link BatchedLogBuffer} decides when a batch seals, {@link
+ * AsyncBatchedLogUploader} decides what an S3 object looks like, and this class decides when
+ * uploads run — the flush triggers (seal on publish, age tick, drain on close), the
+ * upload-concurrency gate, and metrics.
+ *
+ * Delivery is at-most-once by design: publish never waits for queue capacity, full buffers drop
+ * new lines, and close drains best-effort within {@code shutdownTimeout}.
+ */
+public final class BillingUploadingLogHandler extends Handler {
+
+ // Logger for this handler, not the destination we are sending events to.
+ private static final Logger LOGGER = LoggerFactory.getLogger(BillingUploadingLogHandler.class);
+
+ /**
+ * When true means the Handler has been closed via {@link #close()} and it will silently drop any
+ * further calls to publish log entries. This also cause the upload thread to empty the buffer
+ */
+ private final AtomicBoolean isClosed = new AtomicBoolean(false);
+
+ /**
+ * Disposable permitting system for forcing wakeup in the uploading thread. startUploading() will
+ * tryAcquire() but because the permit count is 0 will always timeout, this is the timeout to wake
+ * and check buffer. When we want to force wakeup, e.g. flush(), we call release() that means any
+ * tryAcquire() returns and decrements count to 0. Resetting back to initial state. Because the
+ * wakeup permit lasts until tryAcquire it removes race conditions that could happen when
+ * flush()/notify() on an object lands before the upload thread is in wait() - if we used
+ * Object.notify() and .wait()
+ */
+ private final Semaphore wakeupPermit = new Semaphore(0);
+
+ /**
+ * There is only 1 permit for the upload process to be runnning. When {@link #startUploading()}
+ * starts it takes the permit, gives it back when the function exits (after {@link #close()}.
+ * close() uses this to make sure uploading has finished.
+ */
+ private final Semaphore uploadPermit = new Semaphore(1);
+
+ private final AsyncBatchedLogUploader uploader;
+ private final BatchedLogBuffer buffer;
+ private final Duration uploadSleepDuration;
+ private final Duration uploaderSafetyDeadline;
+ private final Duration uploadShutdownDeadline;
+
+ /** See {@link BillingS3HandlerInstaller} */
+ BillingUploadingLogHandler(
+ BatchedLogBuffer buffer,
+ AsyncBatchedLogUploader uploader,
+ Duration uploadSleepDuration,
+ Duration uploaderSafetyDeadline,
+ Duration uploadShutdownDeadline) {
+
+ this.buffer = Objects.requireNonNull(buffer, "buffer must not be null");
+ this.uploader = Objects.requireNonNull(uploader, "uploader must not be null");
+ this.uploadSleepDuration =
+ Objects.requireNonNull(uploadSleepDuration, "uploadSleepDuration must not be null");
+ this.uploaderSafetyDeadline =
+ Objects.requireNonNull(uploaderSafetyDeadline, "uploaderSafetyDeadline must not be null");
+ this.uploadShutdownDeadline =
+ Objects.requireNonNull(uploadShutdownDeadline, "uploadShutdownDeadline must not be null");
+ }
+
+ /**
+ * WARNING - sets the flag for closing but does not run the full close. just here for testing how
+ * uploading wakes up when flush called.
+ */
+ @VisibleForTesting
+ void unsafeClose() {
+ LOGGER.warn("WARNING - unsafeClose() called, must only be used in testing");
+ isClosed.set(true);
+ }
+
+ /**
+ * WARNING - acquires the upload permit, this stops the startUpload() function and close() from
+ * working normally. For testing only.
+ */
+ @VisibleForTesting
+ void unsafeAcquireUploadPermit() {
+ LOGGER.warn("WARNING - unsafeAcquireUploadPermit() called, must only be used in testing");
+ uploadPermit.acquireUninterruptibly();
+ }
+
+ // ============================================================
+ // Overrides for java.util.logging.Handler
+ // ============================================================
+
+ /**
+ * Buffers and then published the record to S3.
+ *
+ * @param record description of the log event. A null record is silently ignored and is not
+ * published
+ */
+ @Override
+ public void publish(LogRecord record) {
+
+ // Sanity check
+ if (record == null) {
+ return;
+ }
+
+ if (isClosed.get()) {
+ LOGGER.warn("publish() - called when closed, dropping record:{}", record);
+ return;
+ }
+
+ // buffer handles metrics
+ if (!buffer.offer(record)) {
+ if (LOGGER.isDebugEnabled()) {
+ LOGGER.debug(
+ "publish() - buffer.offer() rejected, dropping record:{}", record.getMessage());
+ }
+ } else if (LOGGER.isTraceEnabled()) {
+ LOGGER.trace("publish() - buffer.offer() accepted, record:{}", record.getMessage());
+ }
+ }
+
+ /**
+ * Wakes up the uploading thread to check the buffer for batches.
+ *
+ * This will only drain the buffer fully (i.e. including partial batches) if {@link #close()}
+ * is called or {@link #isClosed} is set.
+ */
+ @Override
+ public void flush() {
+ maybeTrace("flush() - called");
+ notifyUploading();
+ }
+
+ /**
+ * Closes the LogHandler so that it will drop any records sent to {@link #publish(LogRecord)} and
+ * drain the buffer fully to send all batches to S3.
+ */
+ @Override
+ public void close() {
+
+ LOGGER.info(
+ "closing() - marking handler closed, flushing, and waiting for uploads to complete. uploadShutdownDeadline:{}",
+ uploadShutdownDeadline);
+
+ // mark as closed to stop accepting further events and tell the upload thread
+ // to drain the buffer fully.
+ isClosed.set(true);
+ flush();
+
+ try {
+ // Check uploading is not running by trying to get the single uploading permit
+ // TODO: move timeout to config
+ if (!uploadPermit.tryAcquire(uploadShutdownDeadline.toMillis(), TimeUnit.MILLISECONDS)) {
+ LOGGER.warn(
+ "close() - Failed to get uploading permit, upload failed to stop. uploadShutdownDeadline:{}",
+ uploadShutdownDeadline);
+ } else {
+ uploadPermit.release();
+ LOGGER.debug("close() - acquired uploading permit, uploading has completed.");
+ }
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ LOGGER.warn("close() - Interrupted waiting for billing upload loop to finish");
+ } finally {
+ uploader.close();
+ }
+ }
+
+ // ============================================================
+ // Flush pipeline
+ // ============================================================
+
+ /**
+ * Call this on a worker thread to start uploading, will start a loop of waiting for batches from
+ * the buffer and uploading them.
+ */
+ void startUploading() {
+
+ LOGGER.info("startUploading() - using buffer:{}, uploader:{}", buffer, uploader);
+
+ boolean hasPermit = false;
+ BatchedLogBuffer.Batch batch;
+ try {
+ if (!(hasPermit = uploadPermit.tryAcquire())) {
+ throw new IllegalStateException(
+ "startUploading() - unable to acquire uploadPermit, was function already called?");
+ }
+
+ while (true) {
+
+ // if the handler is closed we do not want to go to sleep again because it is closing
+ // down.
+ if (!isClosed.get()) {
+ try {
+ // waiting will release the synchronized monitor
+ maybeTrace(
+ "startUploading() - waiting for wakeupPermit. isClosed:{}, uploadSleepDuration:{}",
+ isClosed.get(),
+ uploadSleepDuration);
+ var acquiredWakePermit =
+ wakeupPermit.tryAcquire(uploadSleepDuration.toMillis(), TimeUnit.MILLISECONDS);
+ // is not important if we got a permit to wake, or timed out, just for logging
+ maybeTrace(
+ "startUploading() - wakeup permit or timeout, isClosed:{}, acquiredWakePermit:{}",
+ isClosed.get(),
+ acquiredWakePermit);
+
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ dumpBuffer();
+ return;
+ }
+ } else {
+ maybeTrace(
+ "startUploading() - not waiting for wakeupPermit because isClosed:{}",
+ isClosed.get());
+ }
+
+ // Get the next batches, if isClosed is true then we want to drain all events
+ // which may mean creating a batch when we do not have a full one.
+ while ((batch = buffer.nextBatch(isClosed.get())) != null) {
+ uploadBatch(batch);
+ }
+
+ if (isClosed.get()) {
+ // Handler is closing down, time to get out of this crazy loop
+ break;
+ }
+ }
+ } finally {
+ // release the uploading permit if we have it, done with the uploading lifestyle
+ if (hasPermit) {
+ uploadPermit.release();
+ maybeTrace("startUploading() - releasing upload permit");
+ } else {
+ maybeTrace("startUploading() - upload permit was not acquired, not releasing");
+ }
+ }
+
+ if (!buffer.isEmpty()) {
+ LOGGER.warn(
+ "startUploading() - finished with abandoned billing events, billingQueue.size():{} ",
+ buffer.size());
+ }
+
+ LOGGER.info(
+ "startUploading() - stopped uploading using buffer:{}, uploader:{}", buffer, uploader);
+ }
+
+ /** Adds a permit to the wakeupPermit so the uploading thread will wakeup and do some work. */
+ private void notifyUploading() {
+ wakeupPermit.release();
+ }
+
+ private static void maybeTrace(String message, Object... args) {
+ if (LOGGER.isTraceEnabled()) {
+ LOGGER.trace(message, args);
+ }
+ }
+
+ /**
+ * Creates a Uni that will upload the provided batch.
+ *
+ * As a deferred Uni it does not do any work until something pulls the item, so the caller (see
+ * startUploading()) starts the work and can decide to wait etc.
+ *
+ * @param batch
+ * @return
+ */
+ private void uploadBatch(BatchedLogBuffer.Batch batch) {
+
+ LOGGER.info(
+ "uploadBatch() - starting to upload. uploaderSafetyDeadline:{}, batch:{}",
+ uploaderSafetyDeadline,
+ batch);
+
+ // while the uploader should take of all the timeout and retry logic
+ // as a client of the uploader adding a safety timeout here incase it breaks
+
+ // using deferred so that an error in upload() before it returns the Uni is then
+ // treated as an error through the Uni pipeline
+ var uploadResult =
+ Uni.createFrom()
+ .deferred(() -> uploader.upload(batch))
+ .ifNoItem()
+ .after(uploaderSafetyDeadline)
+ .fail()
+ .onFailure()
+ .recoverWithItem(
+ t -> onUploaderFailure(batch, t)) // TimeoutException id deadline exceeded
+ .await()
+ .indefinitely(); // the deadline above covers it
+
+ if (uploadResult.throwable() == null) {
+ onBatchSuccess(uploadResult);
+ } else {
+ onBatchFailure(uploadResult);
+ }
+ }
+
+ /**
+ * There was an unhandled error from the uploader().
+ *
+ * Could be from in upload() before it returned or from running the Uni to do the upload. Just
+ * map this unhandled back into the UploadResult so we can deal with error in standard way
+ */
+ private AsyncBatchedLogUploader.UploadResult onUploaderFailure(
+ BatchedLogBuffer.Batch batch, Throwable throwable) {
+ LOGGER.error(
+ "onUploaderFailure() - throwable from uploader, adding to UploadResult. batch:{}, throwable:{}",
+ batch,
+ throwable.toString());
+ return new AsyncBatchedLogUploader.UploadResult(batch, throwable);
+ }
+
+ private void onBatchSuccess(AsyncBatchedLogUploader.UploadResult uploadResult) {
+ LOGGER.info("onBatchSuccess() - successfully uploaded batch:{}", uploadResult.batch());
+ }
+
+ private void onBatchFailure(AsyncBatchedLogUploader.UploadResult uploadResult) {
+ LOGGER.error("onBatchFailure() - failed to upload batch:{}", uploadResult.batch());
+ }
+
+ /** TODO: dump buffer or a failed batch to regular logs or whatever */
+ private void dumpBuffer() {}
+
+ private void dumpBatch() {}
+}
diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/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..2870f5f5d7
--- /dev/null
+++ b/src/main/java/io/stargate/sgv2/jsonapi/service/billing/S3BatchedLogUploader.java
@@ -0,0 +1,194 @@
+package io.stargate.sgv2.jsonapi.service.billing;
+
+import static io.stargate.sgv2.jsonapi.util.ClassUtils.classSimpleName;
+
+import io.smallrye.mutiny.Uni;
+import io.stargate.sgv2.jsonapi.metrics.BatchedLogUploaderMetrics;
+import java.net.URI;
+import java.nio.charset.StandardCharsets;
+import java.time.Duration;
+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.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";
+ private 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 BatchedLogUploaderMetrics metrics;
+
+ private S3BatchedLogUploader(
+ S3AsyncClient client, String region, String bucket, BatchedLogUploaderMetrics metrics) {
+ this.client = client;
+ this.region = region;
+ this.bucket = bucket;
+ this.metrics = metrics;
+ }
+
+ /**
+ * Creates a new instance
+ *
+ * @param region
+ * @param bucket
+ * @param endpointOverride
+ * @return
+ */
+ public static S3BatchedLogUploader create(
+ String region, String bucket, String endpointOverride, BatchedLogUploaderMetrics metrics) {
+
+ if (region == null || region.isBlank()) {
+ throw new IllegalArgumentException("region must not be null or blank");
+ }
+ if (bucket == null || bucket.isBlank()) {
+ throw new IllegalArgumentException("bucket must not be null or blank");
+ }
+ Objects.requireNonNull(metrics, "metrics must not be null");
+
+ // Credentials resolve from the SDK's default provider chain (env vars, web-identity/OIDC
+ // token, instance/container roles), left implicit so the client owns — and closes — the
+ // provider. This transparently supports federated (AssumeRoleWithWebIdentity) and
+ // cross-account access: the bucket may live in a different account (per IAM + bucket
+ // policy); its region is set via .region().
+ var builder =
+ S3AsyncClient.builder()
+ .region(Region.of(region))
+ .overrideConfiguration(
+ config ->
+ config
+ .apiCallAttemptTimeout(API_CALL_ATTEMPT_TIMEOUT)
+ .apiCallTimeout(API_CALL_TIMEOUT));
+
+ // Real AWS S3 needs no endpoint: the SDK endpoint rules (s3 SDK's DefaultS3EndpointProvider)
+ // derive https:// {@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 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 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 NOTE: ttest takes 7 or 8 seconds, if you change the sleep time it may mean there are
+ * no batches collected after shutdown because producers go fast
+ */
+ @Test
+ public void multiThreadedProducerConsumer() {
+
+ var fixture = defaultBufferFixture(false);
+
+ // Setup a Consumer thread, it will keep running until we set consumerShutdown
+ var normalBatches = new ArrayList See {@link #defaultBufferFixture(boolean)}
+ */
+ record Fixture(
+ int maxBatchSize,
+ long maxBytes,
+ Duration maxAge,
+ int queueCapacity,
+ List ...
+ */
+ record Slice(int from, int to) {
+
+ public ...
+ */
+ record BufferSnapshot(
+ boolean isEmpty, int size, long queuedBytes, int remainingCapacity, Fixture fixture) {
+
+ 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 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 Mapnull if there is no next batch.
+ */
+ public Batch nextBatch(boolean drainFully) {
+
+ var batchReason = decideNextBatch(drainFully);
+ if (batchReason == null) {
+ return null;
+ }
+
+ List