diff --git a/braintrust-sdk/src/main/java/dev/braintrust/config/BraintrustConfig.java b/braintrust-sdk/src/main/java/dev/braintrust/config/BraintrustConfig.java index 5b05962d..aec2b0ca 100644 --- a/braintrust-sdk/src/main/java/dev/braintrust/config/BraintrustConfig.java +++ b/braintrust-sdk/src/main/java/dev/braintrust/config/BraintrustConfig.java @@ -58,6 +58,30 @@ public final class BraintrustConfig extends BaseConfig { private final Boolean autoConvertAIAttachments = getConfig("BRAINTRUST_AUTO_CONVERT_AI_ATTACHMENTS", true); + /** Maximum number of attachment uploads waiting for the background uploader. */ + private final int attachmentUploaderQueueSize = + assertPositive(getConfig("BRAINTRUST_ATTACHMENT_UPLOADER_QUEUE_SIZE", 1024)); + + /** Per-request timeout for attachment upload HTTP calls. */ + private final Duration attachmentUploaderRequestTimeout = + Duration.ofMillis( + assertPositive( + getConfig( + "BRAINTRUST_ATTACHMENT_UPLOADER_REQUEST_TIMEOUT_MILLIS", + 60_000))); + + /** Maximum number of retries for transient attachment upload failures. */ + private final int attachmentUploaderMaxRetries = + assertPositive(getConfig("BRAINTRUST_ATTACHMENT_UPLOADER_MAX_RETRIES", 8)); + + /** Initial attachment upload retry delay. The uploader doubles it after each failure. */ + private final Duration attachmentUploaderInitialRetryDelay = + Duration.ofMillis( + assertPositive( + getConfig( + "BRAINTRUST_ATTACHMENT_UPLOADER_INITIAL_RETRY_DELAY_MILLIS", + 500))); + /** Custom SSL context for OTLP exporter. Builder-only field, not backed by envars. */ private final SSLContext sslContext; @@ -288,6 +312,32 @@ public Builder autoConvertAIAttachments(boolean value) { return this; } + public Builder attachmentUploaderQueueSize(int queueSize) { + envOverrides.put( + "BRAINTRUST_ATTACHMENT_UPLOADER_QUEUE_SIZE", String.valueOf(queueSize)); + return this; + } + + public Builder attachmentUploaderRequestTimeout(Duration requestTimeout) { + envOverrides.put( + "BRAINTRUST_ATTACHMENT_UPLOADER_REQUEST_TIMEOUT_MILLIS", + String.valueOf(requestTimeout.toMillis())); + return this; + } + + public Builder attachmentUploaderMaxRetries(int maxRetries) { + envOverrides.put( + "BRAINTRUST_ATTACHMENT_UPLOADER_MAX_RETRIES", String.valueOf(maxRetries)); + return this; + } + + public Builder attachmentUploaderInitialRetryDelay(Duration initialRetryDelay) { + envOverrides.put( + "BRAINTRUST_ATTACHMENT_UPLOADER_INITIAL_RETRY_DELAY_MILLIS", + String.valueOf(initialRetryDelay.toMillis())); + return this; + } + public Builder sslContext(SSLContext value) { this.sslContext = value; return this; diff --git a/braintrust-sdk/src/main/java/dev/braintrust/trace/AttachmentProcessor.java b/braintrust-sdk/src/main/java/dev/braintrust/trace/AttachmentProcessor.java index 0f624ff2..2afa9745 100644 --- a/braintrust-sdk/src/main/java/dev/braintrust/trace/AttachmentProcessor.java +++ b/braintrust-sdk/src/main/java/dev/braintrust/trace/AttachmentProcessor.java @@ -140,7 +140,7 @@ private static Pattern buildHeuristic() { String processAndUpload(String json) { if ((!config.autoConvertAIAttachments()) || json == null - || uploader.isShutdown() + || !uploader.isAcceptingJobs() || !BASE64_HEURISTIC.matcher(json).find()) { return json; } diff --git a/braintrust-sdk/src/main/java/dev/braintrust/trace/AttachmentUploader.java b/braintrust-sdk/src/main/java/dev/braintrust/trace/AttachmentUploader.java index 9c7aa7d3..84103fa2 100644 --- a/braintrust-sdk/src/main/java/dev/braintrust/trace/AttachmentUploader.java +++ b/braintrust-sdk/src/main/java/dev/braintrust/trace/AttachmentUploader.java @@ -2,6 +2,7 @@ import com.fasterxml.jackson.annotation.JsonProperty; import dev.braintrust.api.BraintrustOpenApiClient; +import dev.braintrust.config.BraintrustConfig; import dev.braintrust.json.BraintrustJsonMapper; import java.io.IOException; import java.net.URI; @@ -27,8 +28,6 @@ interface AttachmentUploader { /** * Enqueues an attachment for upload. * - *

NOTE: if the upload queue is full, this method will block until space becomes available - * * @param reference the attachment reference metadata * @param data the attachment data to upload * @return true if the attachment was successfully enqueued for upload. False if the uploader @@ -49,7 +48,8 @@ default void forceFlush() { * first, then flush. * * @param timeout the maximum time to wait - * @return true if all uploads completed, false if timed out + * @return true if all uploads accepted through the snapshot succeeded; false on timeout, + * unfinished work at shutdown, or any permanently dropped upload through that snapshot */ boolean forceFlush(@Nonnull Duration timeout); @@ -66,7 +66,8 @@ default void shutdown() { */ void shutdown(@Nonnull Duration timeout); - boolean isShutdown(); + /** Whether uploads can currently be accepted, independent of available queue capacity. */ + boolean isAcceptingJobs(); /** * Background uploader for Braintrust attachments that uploads to S3 via signed URLs. @@ -76,96 +77,94 @@ default void shutdown() { *

    *
  1. Requests a signed upload URL from the Braintrust API *
  2. Uploads the data to the signed URL - *
  3. Reports the upload status (done/error) to the Braintrust API + *
  4. Reports successful uploads to the Braintrust API (best effort) *
* - *

The uploader starts lazily on first enqueue and can be shut down gracefully. + *

The uploader starts lazily on first enqueue and can be shut down gracefully. HTTP 400, + * 413, and 415 from signed-URL requests or object-store uploads drop that attachment with a + * debug log and a best-effort error status report, then continue queued work. Other failures + * retain their original key and bytes ahead of queued work. Admission pauses so new attachments + * remain inline until the retained upload succeeds or is permanently dropped. Recovery retries + * obtain fresh signed URLs and wait exponentially, starting at the configured initial retry + * delay and capped at two minutes. Existing finite HTTP retries and request durations are + * additional to this wait. + * + *

Retryable failures continue indefinitely while running. Explicit shutdown permanently + * closes admission and only retries within its grace period; retained uploads are in-memory and + * may be lost on forced shutdown. */ @Slf4j class S3AttachmentUploader implements AttachmentUploader { - private static final int QUEUE_SIZE = 1024; - /** Default per-request timeout for HTTP calls. */ - private static final Duration DEFAULT_REQUEST_TIMEOUT = Duration.ofSeconds(60); + private static final long MAX_RECOVERY_RETRY_DELAY_MILLIS = 120_000L; - /** Default maximum number of retry attempts for transient failures. */ - private static final int DEFAULT_MAX_RETRIES = 8; - - /** Default initial backoff delay between retries. Doubles on each subsequent attempt. */ - private static final Duration DEFAULT_INITIAL_RETRY_DELAY = Duration.ofMillis(500); + @FunctionalInterface + interface RecoverySleeper { + void sleep(long delayMillis) throws InterruptedException; + } private final BraintrustOpenApiClient apiClient; private final Duration requestTimeout; private final int maxRetries; private final Duration initialRetryDelay; + private final RecoverySleeper recoverySleeper; private final LinkedBlockingQueue queue; - private final AtomicReference worker = new AtomicReference<>(); private final AtomicReference orgId = new AtomicReference<>(); // non thread safe fields must be checked and read under the lock private final Object lock = new Object(); - private boolean rejectNewJobs = false; + private ExecutorService worker; + private boolean shutdownRequested = false; + private boolean uploadsPaused = false; private boolean workerDone = false; - private CountDownLatch currentBatch = new CountDownLatch(1); + private long acceptedJobs = 0; + private long processedJobs = 0; + // FIFO sequence of the first permanently dropped upload; zero means none. + private long firstDroppedJob = 0; /** - * Creates a new attachment uploader with default retry settings. + * Creates a new attachment uploader. * * @param apiClient the Braintrust API client (provides auth, base URL, and HTTP transport) + * @param config attachment uploader settings */ - S3AttachmentUploader(@Nonnull BraintrustOpenApiClient apiClient) { - this( - apiClient, - DEFAULT_REQUEST_TIMEOUT, - DEFAULT_MAX_RETRIES, - DEFAULT_INITIAL_RETRY_DELAY); + S3AttachmentUploader( + @Nonnull BraintrustOpenApiClient apiClient, @Nonnull BraintrustConfig config) { + this(apiClient, config, Thread::sleep); } - /** - * Creates a new attachment uploader with custom retry settings. - * - * @param apiClient the Braintrust API client (provides auth, base URL, and HTTP transport) - * @param requestTimeout the per-request timeout for HTTP calls - * @param maxRetries the maximum number of retry attempts for transient failures - * @param initialRetryDelay the initial backoff delay between retries (doubles on each - * attempt) - */ S3AttachmentUploader( @Nonnull BraintrustOpenApiClient apiClient, - @Nonnull Duration requestTimeout, - int maxRetries, - @Nonnull Duration initialRetryDelay) { - if (requestTimeout.toMillis() < 0) { - throw new IllegalArgumentException("requestTimeout must be >= 0"); - } - if (maxRetries <= 0) { - throw new IllegalArgumentException("maxRetries must be > 0"); - } - if (initialRetryDelay.toMillis() < 0) { - throw new IllegalArgumentException("initialRetryDelay must be >= 0"); - } + @Nonnull BraintrustConfig config, + @Nonnull RecoverySleeper recoverySleeper) { this.apiClient = apiClient; - this.requestTimeout = requestTimeout; - this.maxRetries = maxRetries; - this.initialRetryDelay = initialRetryDelay; - this.queue = new LinkedBlockingQueue<>(QUEUE_SIZE); + this.requestTimeout = config.attachmentUploaderRequestTimeout(); + this.maxRetries = config.attachmentUploaderMaxRetries(); + this.initialRetryDelay = config.attachmentUploaderInitialRetryDelay(); + this.recoverySleeper = recoverySleeper; + this.queue = new LinkedBlockingQueue<>(config.attachmentUploaderQueueSize()); BraintrustShutdownHook.addShutdownHook( BraintrustShutdownHook.ShutdownOrder.ATTACHMENT_UPLOADER, this::shutdown); } @Override public boolean enqueue(@Nonnull AttachmentReference reference, @Nonnull byte[] data) { - if (checkRejectNewJobsThreadSafe()) { - return false; - } try { - ensureWorkerStarted(); - UploadJob job = new UploadJob(reference, data); - return queue.offer(job, 0, TimeUnit.MILLISECONDS); + synchronized (lock) { + if (!isAcceptingJobs()) { + return false; + } + ensureWorkerStarted(); + if (!queue.offer(new UploadJob(reference, data))) { + return false; + } + acceptedJobs++; + return true; + } } catch (Exception e) { log.error("failed to enqueue attachment", e); - shutdown(); + shutdown(Duration.ZERO); return false; } } @@ -173,108 +172,180 @@ public boolean enqueue(@Nonnull AttachmentReference reference, @Nonnull byte[] d @Override @SneakyThrows public boolean forceFlush(@Nonnull Duration timeout) { - return awaitCurrentBatch(timeout.toMillis(), TimeUnit.MILLISECONDS); + long timeoutNanos = timeout.toNanos(); + long started = System.nanoTime(); + synchronized (lock) { + long target = acceptedJobs; + while (processedJobs < target) { + long remaining = timeoutNanos - (System.nanoTime() - started); + if (workerDone || remaining <= 0) { + return false; + } + TimeUnit.NANOSECONDS.timedWait(lock, remaining); + } + return firstDroppedJob == 0 || firstDroppedJob > target; + } } @Override @SneakyThrows public void shutdown(@Nonnull Duration timeout) { + ExecutorService executor; synchronized (lock) { - rejectNewJobs = true; - if (workerDone) { + shutdownRequested = true; + executor = worker; + if (executor == null) { + workerDone = true; + lock.notifyAll(); return; } } - ExecutorService executor = worker.getAndSet(null); - if (executor == null) { - return; - } executor.shutdown(); - if (!executor.awaitTermination(timeout.toMillis(), TimeUnit.MILLISECONDS)) { - log.warn("failed to gracefully shut down s3 upload worker"); + try { + if (!executor.awaitTermination(timeout.toNanos(), TimeUnit.NANOSECONDS)) { + log.warn("failed to gracefully shut down s3 upload worker"); + executor.shutdownNow(); + } + } catch (InterruptedException e) { executor.shutdownNow(); + Thread.currentThread().interrupt(); + throw e; } } @Override - public boolean isShutdown() { - return checkRejectNewJobsThreadSafe(); + public boolean isAcceptingJobs() { + synchronized (lock) { + return !shutdownRequested && !uploadsPaused && !workerDone; + } } // ── Worker lifecycle ────────────────────────────────────────────── - /** - * start worker thread or do nothing if already started - * - *

calling this method does not require the lock - */ + /** Starts the single worker lazily. Must be called under {@link #lock}. */ private void ensureWorkerStarted() { - if (worker.get() == null) { - var newWorker = + if (worker == null) { + worker = Executors.newSingleThreadExecutor( r -> { Thread t = new Thread(r, "braintrust-attachment-uploader"); t.setDaemon(true); return t; }); - if (worker.compareAndSet(null, newWorker)) { - // NOTE: if shutdown is called concurrently job submission may throw an - // exception. This is fine. - newWorker.submit(this::workerLoop); - } else { - // tried to start the worker concurrently. This is fine, we'll just shut down - // and dereference the redundant worker - newWorker.shutdown(); - } + worker.submit(this::workerLoop); } } private void workerLoop() { log.debug("Attachment uploader worker started"); - while ((!checkRejectNewJobsThreadSafe()) || queue.peek() != null) { - UploadJob job = null; - try { - job = queue.poll(100, TimeUnit.MILLISECONDS); + UploadJob job = null; + long recoveryDelay = + Math.min(initialRetryDelay.toMillis(), MAX_RECOVERY_RETRY_DELAY_MILLIS); + try { + while (!Thread.currentThread().isInterrupted()) { + synchronized (lock) { + if (shutdownRequested && job == null && queue.isEmpty()) { + break; + } + } if (job == null) { - finishCurrentBatch(); - } else { - upload(job); + job = queue.poll(100, TimeUnit.MILLISECONDS); + if (job == null) { + continue; + } } - } catch (InterruptedException e) { - // worker thread shutdownNow was invoked - if (!queue.isEmpty()) { - log.warn( - "s3 uploader force shutdown was reached. Dropping {} uploads", - queue.size(), - e); + boolean dropped = false; + try { + upload(job); + } catch (InterruptedException e) { + throw e; + } catch (Exception e) { + // Login wraps checked exceptions; interruption must still terminate + // recovery. + for (Throwable cause = e; cause != null; cause = cause.getCause()) { + if (cause instanceof InterruptedException interrupted) { + throw interrupted; + } + } + if (e instanceof UploadHttpException failure && failure.isPermanent()) { + log.debug( + "Dropping invalid attachment after HTTP {}. key={}", + failure.statusCode, + job.reference().key(), + e); + reportStatus(job.reference().key(), "error", e.getMessage()); + dropped = true; + } else { + boolean firstFailure; + long pending; + synchronized (lock) { + firstFailure = !uploadsPaused; + uploadsPaused = true; + pending = acceptedJobs - processedJobs; + } + log.warn( + firstFailure + ? "Attachment uploads paused after failure; new" + + " attachments will remain inline. key={}" + + " pending={} retryDelayMillis={}" + : "Attachment upload recovery failed; retrying. key={}" + + " pending={} retryDelayMillis={}", + job.reference().key(), + pending, + recoveryDelay, + e); + recoverySleeper.sleep(recoveryDelay); + recoveryDelay = + Math.min(recoveryDelay * 2, MAX_RECOVERY_RETRY_DELAY_MILLIS); + continue; + } } - break; - } catch (Exception e) { - // this only user of this util is our span processor so we'll just fall back to - // sending attachments in span data if an error occurs synchronized (lock) { - rejectNewJobs = true; + processedJobs++; + if (dropped && firstDroppedJob == 0) { + firstDroppedJob = processedJobs; + } + if (uploadsPaused && !shutdownRequested && !dropped) { + log.info( + "Attachment uploads recovered; accepting new attachments." + + " key={}", + job.reference().key()); + } + uploadsPaused = false; + lock.notifyAll(); } - if (job == null) { - log.warn("Failed to upload attachment", e); - } else { - log.warn("Failed to upload attachment key={}", job.reference().key(), e); - reportStatus(job.reference().key(), "error", e.getMessage()); + recoveryDelay = + Math.min(initialRetryDelay.toMillis(), MAX_RECOVERY_RETRY_DELAY_MILLIS); + job = null; + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } finally { + synchronized (lock) { + workerDone = true; + long unresolved = acceptedJobs - processedJobs; + if (unresolved > 0) { + log.warn( + "Attachment uploader stopped with {} unresolved uploads", + unresolved); } - // NOTE: we'll continue the loop attempting uploads of the remaining jobs until - // the queue is drained + queue.clear(); + lock.notifyAll(); } - } - synchronized (lock) { - workerDone = true; - finishCurrentBatch(); log.debug("Attachment uploader worker stopped"); } } - private boolean checkRejectNewJobsThreadSafe() { - synchronized (lock) { - return rejectNewJobs; + private static final class UploadHttpException extends IOException { + private final int statusCode; + + UploadHttpException(int statusCode, String message) { + super(message); + this.statusCode = statusCode; + } + + boolean isPermanent() { + return statusCode == 400 || statusCode == 413 || statusCode == 415; } } @@ -294,6 +365,9 @@ private void upload(@Nonnull UploadJob job) throws IOException, InterruptedExcep job.reference().contentType(), job.data()); + if (Thread.currentThread().isInterrupted()) { + throw new InterruptedException(); + } reportStatus(job.reference().key(), "done", null); } @@ -326,6 +400,10 @@ private void reportStatus( statusMap.put("error_message", errorMessage); } updateUploadStatus(getOrgId(), key, statusMap); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + log.warn( + "Interrupted reporting attachment status key={} status={}", key, status, e); } catch (Exception e) { log.warn("Failed to report attachment status key={} status={}", key, status, e); } @@ -372,7 +450,8 @@ UploadUrlResponse requestUploadUrl( sendWithRetry(requestBuilder.build(), HttpResponse.BodyHandlers.ofString()); if (!isSuccessStatus(response.statusCode())) { - throw new IOException( + throw new UploadHttpException( + response.statusCode(), "Failed to request upload URL: HTTP " + response.statusCode() + " - " @@ -426,7 +505,8 @@ void uploadToSignedUrl( sendWithRetry(requestBuilder.build(), HttpResponse.BodyHandlers.ofString()); if (!isSuccessStatus(response.statusCode())) { - throw new IOException( + throw new UploadHttpException( + response.statusCode(), "Failed to upload to object store: HTTP " + response.statusCode() + " - " @@ -565,24 +645,6 @@ private static void addAzureBlobHeaders( } } - // ── Batch coordination ──────────────────────────────────────────── - - private void finishCurrentBatch() { - synchronized (lock) { - currentBatch.countDown(); - currentBatch = new CountDownLatch(1); - } - } - - private boolean awaitCurrentBatch(long timeout, TimeUnit timeUnit) - throws InterruptedException { - CountDownLatch latch; - synchronized (lock) { - latch = currentBatch; - } - return latch.await(timeout, timeUnit); - } - // ── DTOs ────────────────────────────────────────────────────────── private record UploadJob(AttachmentReference reference, byte[] data) {} diff --git a/braintrust-sdk/src/main/java/dev/braintrust/trace/BraintrustSpanProcessor.java b/braintrust-sdk/src/main/java/dev/braintrust/trace/BraintrustSpanProcessor.java index c243bb36..4d5d7f53 100644 --- a/braintrust-sdk/src/main/java/dev/braintrust/trace/BraintrustSpanProcessor.java +++ b/braintrust-sdk/src/main/java/dev/braintrust/trace/BraintrustSpanProcessor.java @@ -54,7 +54,7 @@ public class BraintrustSpanProcessor implements SpanProcessor { new AttachmentProcessor( config, new AttachmentUploader.S3AttachmentUploader( - BraintrustOpenApiClient.of(config))); + BraintrustOpenApiClient.of(config), config)); } private static List buildSamplers(BraintrustConfig config) { diff --git a/braintrust-sdk/src/test/java/dev/braintrust/config/BraintrustConfigTest.java b/braintrust-sdk/src/test/java/dev/braintrust/config/BraintrustConfigTest.java index bf7a9b8a..27847215 100644 --- a/braintrust-sdk/src/test/java/dev/braintrust/config/BraintrustConfigTest.java +++ b/braintrust-sdk/src/test/java/dev/braintrust/config/BraintrustConfigTest.java @@ -4,6 +4,7 @@ import java.lang.reflect.Field; import java.lang.reflect.Method; +import java.time.Duration; import java.util.Arrays; import java.util.List; import java.util.Set; @@ -74,6 +75,27 @@ public void testBuilderHasMethodForEveryField() { } } + @Test + void attachmentUploaderSettingsKeepDefaultsAndAllowOverrides() { + var defaults = BraintrustConfig.builder().build(); + assertEquals(1024, defaults.attachmentUploaderQueueSize()); + assertEquals(Duration.ofSeconds(60), defaults.attachmentUploaderRequestTimeout()); + assertEquals(8, defaults.attachmentUploaderMaxRetries()); + assertEquals(Duration.ofMillis(500), defaults.attachmentUploaderInitialRetryDelay()); + + var configured = + BraintrustConfig.builder() + .attachmentUploaderQueueSize(64) + .attachmentUploaderRequestTimeout(Duration.ofSeconds(5)) + .attachmentUploaderMaxRetries(2) + .attachmentUploaderInitialRetryDelay(Duration.ofMillis(25)) + .build(); + assertEquals(64, configured.attachmentUploaderQueueSize()); + assertEquals(Duration.ofSeconds(5), configured.attachmentUploaderRequestTimeout()); + assertEquals(2, configured.attachmentUploaderMaxRetries()); + assertEquals(Duration.ofMillis(25), configured.attachmentUploaderInitialRetryDelay()); + } + @Test void rejectsOtelExportBatchSizeLargerThanQueue() { var thrown = diff --git a/braintrust-sdk/src/test/java/dev/braintrust/trace/AttachmentUploaderTest.java b/braintrust-sdk/src/test/java/dev/braintrust/trace/AttachmentUploaderTest.java index 5863d39f..c31fa380 100644 --- a/braintrust-sdk/src/test/java/dev/braintrust/trace/AttachmentUploaderTest.java +++ b/braintrust-sdk/src/test/java/dev/braintrust/trace/AttachmentUploaderTest.java @@ -1,37 +1,75 @@ package dev.braintrust.trace; import static com.github.tomakehurst.wiremock.client.WireMock.*; +import static com.github.tomakehurst.wiremock.core.WireMockConfiguration.wireMockConfig; +import static java.nio.charset.StandardCharsets.UTF_8; import static org.junit.jupiter.api.Assertions.*; -import com.github.tomakehurst.wiremock.junit5.WireMockRuntimeInfo; -import com.github.tomakehurst.wiremock.junit5.WireMockTest; +import com.github.tomakehurst.wiremock.extension.ResponseDefinitionTransformerV2; +import com.github.tomakehurst.wiremock.http.ResponseDefinition; +import com.github.tomakehurst.wiremock.junit5.WireMockExtension; +import com.github.tomakehurst.wiremock.stubbing.ServeEvent; import dev.braintrust.api.BraintrustOpenApiClient; import dev.braintrust.config.BraintrustConfig; +import dev.braintrust.json.BraintrustJsonMapper; import java.time.Duration; +import java.util.ArrayList; +import java.util.Base64; +import java.util.HashMap; import java.util.Map; +import java.util.concurrent.BlockingQueue; +import java.util.concurrent.Callable; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.FutureTask; +import java.util.concurrent.LinkedBlockingQueue; +import java.util.concurrent.Semaphore; +import java.util.concurrent.TimeUnit; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Nested; import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.RegisterExtension; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.CsvSource; +import org.junit.jupiter.params.provider.ValueSource; -@WireMockTest public class AttachmentUploaderTest { + private static final Duration WAIT = Duration.ofSeconds(5); + private static final ResponseGates RESPONSE_GATES = new ResponseGates(); + + @RegisterExtension + static WireMockExtension wireMock = + WireMockExtension.newInstance() + .options(wireMockConfig().dynamicPort().extensions(RESPONSE_GATES)) + .configureStaticDsl(true) + .build(); + private AttachmentUploader.S3AttachmentUploader uploader; private String baseUrl; + private BraintrustConfig config; @BeforeEach - void setUp(WireMockRuntimeInfo wmRuntimeInfo) { - baseUrl = wmRuntimeInfo.getHttpBaseUrl(); - var config = BraintrustConfig.builder().apiKey("test-api-key").apiUrl(baseUrl).build(); + void setUp() { + baseUrl = wireMock.getRuntimeInfo().getHttpBaseUrl(); + config = + BraintrustConfig.builder() + .apiKey("test-api-key") + .apiUrl(baseUrl) + .attachmentUploaderRequestTimeout(Duration.ofMillis(10_000)) + .attachmentUploaderMaxRetries(1) + .attachmentUploaderInitialRetryDelay(Duration.ofMillis(50)) + .build(); var apiClient = BraintrustOpenApiClient.of(config); - uploader = - new AttachmentUploader.S3AttachmentUploader( - apiClient, Duration.ofMillis(10_000), 1, Duration.ofMillis(50)); + uploader = new AttachmentUploader.S3AttachmentUploader(apiClient, config); } @AfterEach void tearDown() { uploader.shutdown(Duration.ofSeconds(0)); + RESPONSE_GATES.reset(); } private void stubLoginAndUploadFlow() { @@ -62,6 +100,100 @@ private void stubLoginAndUploadFlow() { // ── Worker / queue integration tests ────────────────────────────── + @Test + void attachmentFreeProcessorTrafficNeverCallsUploadEndpoints() { + var processor = new AttachmentProcessor(config, uploader); + String text = "{\"action\":{\"query\":\"weather\"},\"status\":\"completed\"}"; + for (int i = 0; i < 10_000; i++) { + assertEquals(text, processor.processAndUpload(text)); + } + assertTrue(uploader.forceFlush(WAIT)); + verify(0, anyRequestedFor(anyUrl())); + } + + @Test + void concurrentProcessorUploadsRecoverAfterSignedUrlFailure() throws Exception { + var sleeper = useControlledRecovery(Duration.ofMillis(50), 128); + stubLoginAndUploadFlow(); + var outage = + stubFor(post(urlEqualTo("/attachment")).willReturn(aResponse().withStatus(403))); + var processor = new AttachmentProcessor(config, uploader); + Map accepted; + // Hold the first failure until every producer has enqueued its initial batch. + try (var response = RESPONSE_GATES.hold("/attachment")) { + accepted = processConcurrently(processor, "before outage"); + response.awaitRequest(); + response.release(); + assertEquals(50L, sleeper.awaitDelay()); + } + assertFalse(uploader.forceFlush(Duration.ZERO)); + var signedUrlRequests = findAll(postRequestedFor(urlEqualTo("/attachment"))); + assertEquals(1, signedUrlRequests.size(), "The failed job must block later queued jobs"); + String retainedKey = + BraintrustJsonMapper.get() + .readTree(signedUrlRequests.get(0).getBody()) + .get("key") + .asText(); + var paused = processConcurrently(processor, "during outage"); + for (var entry : paused.entrySet()) { + assertEquals(dataUriJson(entry.getKey().getBytes(UTF_8)), entry.getValue()); + } + verify(1, postRequestedFor(urlEqualTo("/attachment"))); + verify(0, putRequestedFor(anyUrl())); + + removeStub(outage); + // Each signed URL identifies its attachment so byte/key mismatches are observable. + stubFor( + post(urlEqualTo("/attachment")) + .willReturn( + aResponse() + .withHeader("Content-Type", "application/json") + .withBody( + "{\"signedUrl\":\"" + + baseUrl + + "/upload/{{jsonPath request.body" + + " '$.key'}}\",\"headers\":{}}") + .withTransformers("response-template"))); + stubFor(put(urlPathMatching("/upload/.*")).willReturn(aResponse().withStatus(200))); + sleeper.release(); + assertTrue( + uploader.forceFlush(WAIT), + "Retained attachments must recover on the same uploader"); + + accepted.putAll(processConcurrently(processor, "after recovery")); + assertTrue(uploader.forceFlush(WAIT)); + var expected = new HashMap(); + for (var entry : accepted.entrySet()) { + var reference = BraintrustJsonMapper.get().readTree(entry.getValue()).get("url"); + assertEquals("braintrust_attachment", reference.get("type").asText()); + assertNull( + expected.put(reference.get("key").asText(), entry.getKey().getBytes(UTF_8)), + "Concurrent attachments must have distinct keys"); + } + var uploads = findAll(putRequestedFor(urlPathMatching("/upload/.*"))); + assertEquals( + expected.size(), uploads.size(), "Every accepted attachment uploads exactly once"); + for (var upload : uploads) { + String key = upload.getUrl().substring("/upload/".length()); + byte[] bytes = expected.remove(key); + assertNotNull(bytes, "Unexpected or duplicate upload: " + key); + assertArrayEquals( + bytes, upload.getBody(), "Bytes must match the exported reference: " + key); + } + assertTrue(expected.isEmpty(), "No accepted attachments may be lost"); + verify(accepted.size() + 1, postRequestedFor(urlEqualTo("/attachment"))); + verify( + 2, + postRequestedFor(urlEqualTo("/attachment")) + .withRequestBody(matchingJsonPath("$.key", equalTo(retainedKey)))); + verify( + accepted.size(), + postRequestedFor(urlEqualTo("/attachment/status")) + .withRequestBody( + matchingJsonPath("$.status.upload_status", equalTo("done")))); + verifyNoErrorStatus(); + } + @Test void enqueueUploadsSuccessfully() throws Exception { stubLoginAndUploadFlow(); @@ -80,26 +212,517 @@ void enqueueUploadsSuccessfully() throws Exception { } @Test - void enqueueRejectsAfterShutdown() { - assertDoesNotThrow(() -> uploader.shutdown()); + void enqueueRejectsAfterShutdownBeforeWorkerStarts() throws Exception { + assertTrue(uploader.forceFlush(Duration.ZERO)); + uploader.shutdown(Duration.ZERO); + assertFalse(uploader.isAcceptingJobs()); + assertTrue(uploader.forceFlush(Duration.ZERO)); var ref = AttachmentReference.create("test.json", "application/json"); assertFalse(uploader.enqueue(ref, "data".getBytes())); } @Test - void uploadFailureShutsDownWorker() throws Exception { - stubFor( - post(urlEqualTo("/api/apikey/login")) - .willReturn( - aResponse() - .withStatus(200) - .withHeader("Content-Type", "application/json") - .withBody( - "{\"org_info\":[{\"id\":\"org-123\",\"name\":\"test-org\"}]}"))); + void retainsFailedAttachmentAndResumesProcessorAfterRecovery() throws Exception { + var sleeper = useControlledRecovery(Duration.ofMillis(500), 2); + stubLoginAndUploadFlow(); + var first = AttachmentReference.create("first.txt", "text/plain"); + var second = AttachmentReference.create("second.txt", "text/plain"); + byte[] firstBytes = "original first attachment".getBytes(UTF_8); + byte[] secondBytes = "queued second attachment".getBytes(UTF_8); + stubSignedUrl(first, "/first"); + stubSignedUrl(second, "/second"); + stubFor(put(urlEqualTo("/first")).willReturn(aResponse().withStatus(403))); + stubFor(put(urlEqualTo("/second")).willReturn(aResponse().withStatus(200))); + var processor = new AttachmentProcessor(config, uploader); + String inline = dataUriJson("paused attachment remains inline".getBytes(UTF_8)); + + try (var firstResponse = RESPONSE_GATES.hold("/first"); + var secondResponse = RESPONSE_GATES.hold("/second")) { + assertTrue(uploader.enqueue(first, firstBytes)); + firstResponse.awaitRequest(); + assertTrue(uploader.enqueue(second, secondBytes)); + firstResponse.release(); + assertEquals(500L, sleeper.awaitDelay()); + + assertFalse(uploader.isAcceptingJobs()); + assertFalse( + uploader.enqueue( + AttachmentReference.create("rejected", "text/plain"), firstBytes)); + assertEquals(inline, processor.processAndUpload(inline)); + verify(0, putRequestedFor(urlEqualTo("/second"))); + assertFalse(uploader.forceFlush(Duration.ofMillis(20))); + verifyNoErrorStatus(); + + // A new URL proves recovery renews the signed URL rather than reusing an expired one. + stubSignedUrl(first, "/first-recovered"); + stubFor(put(urlEqualTo("/first-recovered")).willReturn(aResponse().withStatus(200))); + sleeper.release(); + secondResponse.awaitRequest(); + assertTrue(uploader.isAcceptingJobs(), "Recovery must not wait for the entire queue"); + verify( + 2, + postRequestedFor(urlEqualTo("/attachment")) + .withRequestBody(matchingJsonPath("$.key", equalTo(first.key())))); + verify( + 1, + putRequestedFor(urlEqualTo("/first")) + .withRequestBody(binaryEqualTo(firstBytes))); + verify( + 1, + putRequestedFor(urlEqualTo("/first-recovered")) + .withRequestBody(binaryEqualTo(firstBytes))); + + byte[] resumedBytes = "new attachment after recovery".getBytes(UTF_8); + var converted = + BraintrustJsonMapper.get() + .readTree(processor.processAndUpload(dataUriJson(resumedBytes))); + var reference = converted.get("url"); + assertEquals("braintrust_attachment", reference.get("type").asText()); + String resumedKey = reference.get("key").asText(); + secondResponse.release(); + assertTrue(uploader.forceFlush(WAIT)); + verify( + 1, + putRequestedFor(urlEqualTo("/second")) + .withRequestBody(binaryEqualTo(secondBytes))); + verify( + 1, + postRequestedFor(urlEqualTo("/attachment")) + .withRequestBody(matchingJsonPath("$.key", equalTo(resumedKey)))); + verify( + 1, + putRequestedFor(urlEqualTo("/upload")) + .withRequestBody(binaryEqualTo(resumedBytes))); + verifyNoErrorStatus(); + } + } + + @ParameterizedTest + @CsvSource({"400, true", "413, true", "415, true", "400, false", "413, false", "415, false"}) + void invalidAttachmentIsDroppedWithoutBlockingQueuedUploads(int status, boolean signedUrl) + throws Exception { + useControlledRecovery(Duration.ofMillis(500), 2); + stubLoginAndUploadFlow(); + var invalid = AttachmentReference.create("invalid", "text/plain"); + var queued = AttachmentReference.create("queued", "text/plain"); + stubSignedUrl(invalid, "/invalid"); + stubSignedUrl(queued, "/queued"); + if (signedUrl) { + stubFor( + post(urlEqualTo("/attachment")) + .withRequestBody(matchingJsonPath("$.key", equalTo(invalid.key()))) + .willReturn(aResponse().withStatus(status))); + } else { + stubFor(put(urlEqualTo("/invalid")).willReturn(aResponse().withStatus(status))); + } + stubFor(put(urlEqualTo("/queued")).willReturn(aResponse().withStatus(200))); + + try (var invalidResponse = RESPONSE_GATES.hold(signedUrl ? "/attachment" : "/invalid"); + var queuedResponse = RESPONSE_GATES.hold("/queued")) { + assertTrue(uploader.enqueue(invalid, "invalid".getBytes(UTF_8))); + invalidResponse.awaitRequest(); + assertTrue(uploader.enqueue(queued, "queued".getBytes(UTF_8))); + invalidResponse.release(); + queuedResponse.awaitRequest(); + assertTrue(uploader.isAcceptingJobs()); + var later = AttachmentReference.create("later", "text/plain"); + assertTrue(uploader.enqueue(later, "later".getBytes(UTF_8))); + queuedResponse.release(); + uploader.shutdown(WAIT); + assertFalse(uploader.forceFlush(Duration.ZERO), "A dropped upload is not successful"); + verify( + 1, + postRequestedFor(urlEqualTo("/attachment")) + .withRequestBody(matchingJsonPath("$.key", equalTo(invalid.key())))); + verify(signedUrl ? 0 : 1, putRequestedFor(urlEqualTo("/invalid"))); + verify( + 1, + postRequestedFor(urlEqualTo("/attachment/status")) + .withRequestBody(matchingJsonPath("$.key", equalTo(invalid.key()))) + .withRequestBody( + matchingJsonPath("$.status.upload_status", equalTo("error"))) + .withRequestBody(matchingJsonPath("$.status.error_message"))); + verify( + 1, + putRequestedFor(urlEqualTo("/queued")) + .withRequestBody(binaryEqualTo("queued".getBytes(UTF_8)))); + verify( + 1, + postRequestedFor(urlEqualTo("/attachment/status")) + .withRequestBody(matchingJsonPath("$.key", equalTo(later.key()))) + .withRequestBody( + matchingJsonPath("$.status.upload_status", equalTo("done")))); + } + } + + @ParameterizedTest + @ValueSource(booleans = {false, true}) + void permanentFailureDuringRecoveryDropsJobAndReopensAdmission(boolean statusReportingFails) + throws Exception { + var sleeper = useControlledRecovery(Duration.ofMillis(500), 2); + stubLoginAndUploadFlow(); + stubFor(put(urlEqualTo("/upload")).willReturn(aResponse().withStatus(403))); + var invalid = AttachmentReference.create("invalid", "text/plain"); + assertTrue(uploader.enqueue(invalid, "invalid".getBytes(UTF_8))); + assertEquals(500L, sleeper.awaitDelay()); + verifyNoErrorStatus(); + if (statusReportingFails) { + stubFor( + post(urlEqualTo("/attachment/status")) + .withRequestBody(matchingJsonPath("$.key", equalTo(invalid.key()))) + .willReturn(aResponse().withStatus(500))); + } + stubFor(put(urlEqualTo("/upload")).willReturn(aResponse().withStatus(415))); + sleeper.release(); + assertFalse(uploader.forceFlush(WAIT)); + assertTrue(uploader.isAcceptingJobs(), "Dropping the retained job must release the pause"); + stubFor(put(urlEqualTo("/upload")).willReturn(aResponse().withStatus(200))); + assertTrue( + uploader.enqueue( + AttachmentReference.create("later", "text/plain"), + "later".getBytes(UTF_8))); + uploader.shutdown(WAIT); + verify( + 2, + putRequestedFor(urlEqualTo("/upload")) + .withRequestBody(binaryEqualTo("invalid".getBytes(UTF_8)))); + verify( + 1, + putRequestedFor(urlEqualTo("/upload")) + .withRequestBody(binaryEqualTo("later".getBytes(UTF_8)))); + verify( + statusReportingFails ? 2 : 1, + postRequestedFor(urlEqualTo("/attachment/status")) + .withRequestBody(matchingJsonPath("$.key", equalTo(invalid.key()))) + .withRequestBody( + matchingJsonPath("$.status.upload_status", equalTo("error"))) + .withRequestBody(matchingJsonPath("$.status.error_message"))); + } + + @Test + void exhaustedHttpRetriesEnterRecoveryAndCanSucceed() throws Exception { + var sleeper = useControlledRecovery(Duration.ofMillis(50), 2); + stubLoginAndUploadFlow(); + stubFor(put(urlEqualTo("/upload")).willReturn(aResponse().withStatus(500))); + var reference = AttachmentReference.create("retry.txt", "text/plain"); + byte[] bytes = "transient upload".getBytes(UTF_8); + assertTrue(uploader.enqueue(reference, bytes)); + + assertEquals(50L, sleeper.awaitDelay()); + verify(2, putRequestedFor(urlEqualTo("/upload"))); + assertFalse(uploader.isAcceptingJobs()); + assertFalse(uploader.forceFlush(Duration.ZERO)); + verifyNoErrorStatus(); + + stubFor(put(urlEqualTo("/upload")).willReturn(aResponse().withStatus(200))); + sleeper.release(); + assertTrue(uploader.forceFlush(WAIT)); + assertTrue(uploader.isAcceptingJobs()); + verify( + 2, + postRequestedFor(urlEqualTo("/attachment")) + .withRequestBody(matchingJsonPath("$.key", equalTo(reference.key())))); + verify(3, putRequestedFor(urlEqualTo("/upload")).withRequestBody(binaryEqualTo(bytes))); + } + + @Test + void recoveryBackoffStaysCappedIndefinitelyAndResetsAfterSuccess() throws Exception { + var sleeper = useControlledRecovery(Duration.ofMillis(500), 2); + stubLoginAndUploadFlow(); + var first = AttachmentReference.create("first.txt", "text/plain"); + var queued = AttachmentReference.create("queued.txt", "text/plain"); + stubSignedUrl(first, "/first"); + stubSignedUrl(queued, "/queued"); + stubFor(put(urlEqualTo("/first")).willReturn(aResponse().withStatus(403))); + stubFor(put(urlEqualTo("/queued")).willReturn(aResponse().withStatus(200))); + try (var firstResponse = RESPONSE_GATES.hold("/first")) { + assertTrue(uploader.enqueue(first, "first".getBytes(UTF_8))); + firstResponse.awaitRequest(); + assertTrue(uploader.enqueue(queued, "queued".getBytes(UTF_8))); + firstResponse.release(); + long[] delays = { + 500, 1_000, 2_000, 4_000, 8_000, 16_000, 32_000, 64_000, 120_000, 120_000 + }; + for (int i = 0; i < delays.length; i++) { + assertEquals(delays[i], sleeper.awaitDelay()); + assertFalse(uploader.isAcceptingJobs()); + verify(i + 1, putRequestedFor(urlEqualTo("/first"))); + verify( + 0, + postRequestedFor(urlEqualTo("/attachment")) + .withRequestBody(matchingJsonPath("$.key", equalTo(queued.key())))); + if (i + 1 < delays.length) { + sleeper.release(); + } + } + stubFor(put(urlEqualTo("/first")).willReturn(aResponse().withStatus(200))); + sleeper.release(); + assertTrue(uploader.forceFlush(WAIT)); + verify(1, putRequestedFor(urlEqualTo("/queued"))); + } + + var later = AttachmentReference.create("later.txt", "text/plain"); + stubFor(put(urlEqualTo("/upload")).willReturn(aResponse().withStatus(403))); + assertTrue(uploader.enqueue(later, "later".getBytes(UTF_8))); + assertEquals(500L, sleeper.awaitDelay(), "A new outage starts at the initial delay"); + stubFor(put(urlEqualTo("/upload")).willReturn(aResponse().withStatus(200))); + sleeper.release(); + assertTrue(uploader.forceFlush(WAIT)); + verifyNoErrorStatus(); + } + + @Test + void zeroTimeoutShutdownInterruptsRecoveryWithoutAnotherAttempt() throws Exception { + var sleeper = useControlledRecovery(Duration.ofMillis(500), 2); + stubLoginAndUploadFlow(); + stubFor(put(urlEqualTo("/upload")).willReturn(aResponse().withStatus(403))); + var reference = AttachmentReference.create("pending.txt", "text/plain"); + var queued = AttachmentReference.create("queued.txt", "text/plain"); + try (var response = RESPONSE_GATES.hold("/upload")) { + assertTrue(uploader.enqueue(reference, "pending".getBytes(UTF_8))); + response.awaitRequest(); + assertTrue(uploader.enqueue(queued, "queued".getBytes(UTF_8))); + response.release(); + assertEquals(500L, sleeper.awaitDelay()); + } + + try (var shutdown = + new WaitingCall<>( + () -> { + uploader.shutdown(Duration.ZERO); + return null; + })) { + shutdown.result(); + } + sleeper.awaitInterruption(); + assertFalse(uploader.isAcceptingJobs()); + assertFalse(uploader.enqueue(reference, "rejected".getBytes(UTF_8))); + assertFalse(uploader.forceFlush(WAIT), "Interrupted work is not completed work"); + uploader.shutdown(Duration.ZERO); + verify(1, postRequestedFor(urlEqualTo("/attachment"))); + verify(1, putRequestedFor(urlEqualTo("/upload"))); + verify(0, postRequestedFor(urlEqualTo("/attachment/status"))); + } + + @Test + void successfulRecoveryDuringGracefulShutdownDoesNotReopenAdmission() throws Exception { + var sleeper = useControlledRecovery(Duration.ofMillis(500), 2); + stubLoginAndUploadFlow(); + stubFor(put(urlEqualTo("/upload")).willReturn(aResponse().withStatus(403))); + var reference = AttachmentReference.create("pending.txt", "text/plain"); + assertTrue(uploader.enqueue(reference, "pending".getBytes(UTF_8))); + assertEquals(500L, sleeper.awaitDelay()); + + stubFor(put(urlEqualTo("/upload")).willReturn(aResponse().withStatus(200))); + try (var response = RESPONSE_GATES.hold("/upload")) { + sleeper.release(); + response.awaitRequest(); + try (var shutdown = + new WaitingCall<>( + () -> { + uploader.shutdown(WAIT); + return null; + })) { + shutdown.awaitWaiting(); + assertFalse(uploader.isAcceptingJobs()); + response.release(); + shutdown.result(); + } + } + assertTrue(uploader.forceFlush(WAIT)); + assertFalse(uploader.isAcceptingJobs()); + assertFalse(uploader.enqueue(reference, "rejected".getBytes(UTF_8))); + verify(2, putRequestedFor(urlEqualTo("/upload"))); + verify( + 1, + postRequestedFor(urlEqualTo("/attachment/status")) + .withRequestBody( + matchingJsonPath("$.status.upload_status", equalTo("done")))); + } + + @Test + void fullQueueRejectsWithoutPausingAndCapacityBecomesUsable() throws Exception { + useControlledRecovery(Duration.ofMillis(500), 1); + stubLoginAndUploadFlow(); + var first = AttachmentReference.create("first", "text/plain"); + var second = AttachmentReference.create("second", "text/plain"); + var third = AttachmentReference.create("third", "text/plain"); + stubSignedUrl(first, "/first"); + stubSignedUrl(second, "/second"); + stubFor(put(urlEqualTo("/first")).willReturn(aResponse().withStatus(200))); + stubFor(put(urlEqualTo("/second")).willReturn(aResponse().withStatus(200))); + try (var firstResponse = RESPONSE_GATES.hold("/first"); + var secondResponse = RESPONSE_GATES.hold("/second")) { + assertTrue(uploader.enqueue(first, "first".getBytes(UTF_8))); + firstResponse.awaitRequest(); + assertTrue(uploader.enqueue(second, "second".getBytes(UTF_8))); + assertFalse(uploader.enqueue(third, "third".getBytes(UTF_8))); + assertTrue(uploader.isAcceptingJobs(), "Capacity rejection is not an upload outage"); + firstResponse.release(); + secondResponse.awaitRequest(); + assertTrue(uploader.enqueue(third, "third".getBytes(UTF_8))); + secondResponse.release(); + assertTrue(uploader.forceFlush(WAIT)); + verify( + 1, + postRequestedFor(urlEqualTo("/attachment")) + .withRequestBody(matchingJsonPath("$.key", equalTo(third.key())))); + verify( + 1, + putRequestedFor(urlEqualTo("/upload")) + .withRequestBody(binaryEqualTo("third".getBytes(UTF_8)))); + } + } + + @Test + void concurrentProcessorOverflowStaysInlineUntilQueueDrains() throws Exception { + useControlledRecovery(Duration.ofMillis(50), 1); + stubLoginAndUploadFlow(); + var processor = new AttachmentProcessor(config, uploader); + var accepted = new HashMap(); + String first = "attachment held in the object store"; + try (var response = RESPONSE_GATES.hold("/upload")) { + accepted.put(first, processor.processAndUpload(dataUriJson(first.getBytes(UTF_8)))); + response.awaitRequest(); + // One upload is active and the one-slot waiting queue is initially empty. + var results = processConcurrently(processor, "queue saturation"); + int inline = 0; + for (var entry : results.entrySet()) { + String original = dataUriJson(entry.getKey().getBytes(UTF_8)); + if (original.equals(entry.getValue())) { + inline++; + } else { + accepted.put(entry.getKey(), entry.getValue()); + } + } + assertEquals(2, accepted.size(), "Only the active upload and one queued job fit"); + assertEquals(results.size() - 1, inline, "Overflow must preserve the original JSON"); + assertTrue(uploader.isAcceptingJobs(), "Queue pressure must not pause the uploader"); + assertFalse(uploader.forceFlush(Duration.ZERO)); + response.release(); + assertTrue(uploader.forceFlush(WAIT)); + } + + String later = "attachment after the queue drains"; + accepted.put(later, processor.processAndUpload(dataUriJson(later.getBytes(UTF_8)))); + assertTrue(uploader.forceFlush(WAIT)); + verify(accepted.size(), putRequestedFor(urlEqualTo("/upload"))); + verify(accepted.size(), postRequestedFor(urlEqualTo("/attachment"))); + for (var entry : accepted.entrySet()) { + var reference = BraintrustJsonMapper.get().readTree(entry.getValue()).get("url"); + assertEquals("braintrust_attachment", reference.get("type").asText()); + verify( + 1, + postRequestedFor(urlEqualTo("/attachment")) + .withRequestBody( + matchingJsonPath( + "$.key", equalTo(reference.get("key").asText())))); + verify( + 1, + putRequestedFor(urlEqualTo("/upload")) + .withRequestBody(binaryEqualTo(entry.getKey().getBytes(UTF_8)))); + } + verify( + accepted.size(), + postRequestedFor(urlEqualTo("/attachment/status")) + .withRequestBody( + matchingJsonPath("$.status.upload_status", equalTo("done")))); + } + + @Test + void flushWaitsOnlyForWorkAcceptedBeforeItsSnapshot() throws Exception { + stubLoginAndUploadFlow(); + assertTrue(uploader.forceFlush(Duration.ZERO)); + var first = AttachmentReference.create("first", "text/plain"); + var later = AttachmentReference.create("later", "text/plain"); + stubSignedUrl(first, "/first"); + stubSignedUrl(later, "/later"); + stubFor(put(urlEqualTo("/first")).willReturn(aResponse().withStatus(200))); + stubFor(put(urlEqualTo("/later")).willReturn(aResponse().withStatus(200))); + try (var firstResponse = RESPONSE_GATES.hold("/first"); + var laterResponse = RESPONSE_GATES.hold("/later")) { + assertTrue(uploader.enqueue(first, "first".getBytes(UTF_8))); + firstResponse.awaitRequest(); + try (var flush = new WaitingCall<>(() -> uploader.forceFlush(WAIT))) { + // The dedicated thread has no other blocking operation before forceFlush's wait. + flush.awaitWaiting(); + assertTrue(uploader.enqueue(later, "later".getBytes(UTF_8))); + firstResponse.release(); + laterResponse.awaitRequest(); + assertTrue(flush.result(), "A later accepted job must not extend the snapshot"); + assertFalse(uploader.forceFlush(Duration.ZERO)); + laterResponse.release(); + assertTrue(uploader.forceFlush(WAIT)); + } + } + } + + private Map processConcurrently(AttachmentProcessor processor, String batch) + throws Exception { + int producers = 4; + var ready = new CountDownLatch(producers); + var start = new CountDownLatch(1); + var pool = Executors.newFixedThreadPool(producers); + try { + var futures = new ArrayList>>(); + for (int producer = 0; producer < producers; producer++) { + int id = producer; + futures.add( + pool.submit( + () -> { + ready.countDown(); + assertTrue(start.await(WAIT.toMillis(), TimeUnit.MILLISECONDS)); + var results = new HashMap(); + for (int i = 0; i < 8; i++) { + String content = + batch + " producer=" + id + " attachment=" + i; + results.put( + content, + processor.processAndUpload( + dataUriJson(content.getBytes(UTF_8)))); + } + return results; + })); + } + assertTrue(ready.await(WAIT.toMillis(), TimeUnit.MILLISECONDS)); + start.countDown(); + var results = new HashMap(); + for (var future : futures) { + results.putAll(future.get(WAIT.toMillis(), TimeUnit.MILLISECONDS)); + } + return results; + } finally { + pool.shutdownNow(); + assertTrue( + pool.awaitTermination(WAIT.toMillis(), TimeUnit.MILLISECONDS), + "Processor caller threads must terminate"); + } + } + private ControlledSleeper useControlledRecovery(Duration initialDelay, int queueSize) { + uploader.shutdown(Duration.ZERO); + config = + BraintrustConfig.builder() + .apiKey("test-api-key") + .apiUrl(baseUrl) + .autoConvertAIAttachments(true) + .attachmentUploaderRequestTimeout(Duration.ofSeconds(10)) + .attachmentUploaderMaxRetries(1) + .attachmentUploaderInitialRetryDelay(initialDelay) + .attachmentUploaderQueueSize(queueSize) + .build(); + var sleeper = new ControlledSleeper(); + uploader = + new AttachmentUploader.S3AttachmentUploader( + BraintrustOpenApiClient.of(config), config, sleeper); + return sleeper; + } + + private void stubSignedUrl(AttachmentReference reference, String path) { stubFor( post(urlEqualTo("/attachment")) + .withRequestBody(matchingJsonPath("$.key", equalTo(reference.key()))) .willReturn( aResponse() .withStatus(200) @@ -107,24 +730,155 @@ void uploadFailureShutsDownWorker() throws Exception { .withBody( "{\"signedUrl\":\"" + baseUrl - + "/upload\",\"headers\":{}}"))); - - stubFor( - put(urlEqualTo("/upload")) - .willReturn(aResponse().withStatus(500).withBody("Failed"))); - - stubFor(post(urlEqualTo("/attachment/status")).willReturn(aResponse().withStatus(200))); + + path + + "\",\"headers\":{}}"))); + } - var ref = AttachmentReference.create("test.json", "application/json"); - assertTrue(uploader.enqueue(ref, "data".getBytes())); - // even errors should notify completion - uploader.forceFlush(Duration.ofSeconds(5)); - assertFalse(uploader.enqueue(ref, "data".getBytes())); + private static String dataUriJson(byte[] bytes) { + return "{\"url\":\"data:text/plain;base64," + + Base64.getEncoder().encodeToString(bytes) + + "\"}"; + } + private static void verifyNoErrorStatus() { verify( + 0, postRequestedFor(urlEqualTo("/attachment/status")) - .withRequestBody(containing("\"upload_status\":\"error\"")) - .withRequestBody(containing("\"error_message\""))); + .withRequestBody( + matchingJsonPath("$.status.upload_status", equalTo("error")))); + } + + private static final class ControlledSleeper + implements AttachmentUploader.S3AttachmentUploader.RecoverySleeper { + private final BlockingQueue delays = new LinkedBlockingQueue<>(); + private final Semaphore retries = new Semaphore(0); + private final CountDownLatch interrupted = new CountDownLatch(1); + + @Override + public void sleep(long delayMillis) throws InterruptedException { + delays.add(delayMillis); + try { + retries.acquire(); + } catch (InterruptedException e) { + interrupted.countDown(); + throw e; + } + } + + long awaitDelay() throws InterruptedException { + Long delay = delays.poll(WAIT.toMillis(), TimeUnit.MILLISECONDS); + assertNotNull(delay, "The worker did not enter recovery sleep"); + return delay; + } + + void release() { + retries.release(); + } + + void awaitInterruption() throws InterruptedException { + assertTrue( + interrupted.await(WAIT.toMillis(), TimeUnit.MILLISECONDS), + "Recovery sleep was not interrupted"); + } + } + + private static final class ResponseGates implements ResponseDefinitionTransformerV2 { + private final Map gates = new ConcurrentHashMap<>(); + + ResponseGate hold(String path) { + var gate = new ResponseGate(); + assertNull(gates.putIfAbsent(path, gate), "A response gate already exists for " + path); + return gate; + } + + @Override + public ResponseDefinition transform(ServeEvent event) { + var gate = gates.remove(event.getRequest().getUrl()); + if (gate != null) { + gate.entered.countDown(); + try { + if (!gate.released.await(10, TimeUnit.SECONDS)) { + throw new IllegalStateException("Response gate was not released"); + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new IllegalStateException("Response gate interrupted", e); + } + } + return event.getResponseDefinition(); + } + + void reset() { + gates.values().forEach(ResponseGate::release); + gates.clear(); + } + + @Override + public boolean applyGlobally() { + return true; + } + + @Override + public String getName() { + return "attachment-response-gates"; + } + } + + private static final class ResponseGate implements AutoCloseable { + private final CountDownLatch entered = new CountDownLatch(1); + private final CountDownLatch released = new CountDownLatch(1); + + void awaitRequest() throws InterruptedException { + assertTrue( + entered.await(WAIT.toMillis(), TimeUnit.MILLISECONDS), + "Expected HTTP request did not arrive"); + } + + void release() { + released.countDown(); + } + + @Override + public void close() { + release(); + } + } + + /** Observes only the public caller thread, never the uploader's private state. */ + private static final class WaitingCall implements AutoCloseable { + private final FutureTask call; + private final Thread thread; + + WaitingCall(Callable action) { + call = new FutureTask<>(action); + thread = new Thread(call, "attachment-uploader-test-caller"); + thread.setDaemon(true); + thread.start(); + } + + void awaitWaiting() { + long start = System.nanoTime(); + while (System.nanoTime() - start < WAIT.toNanos()) { + if (thread.getState() == Thread.State.TIMED_WAITING + || thread.getState() == Thread.State.WAITING) { + return; + } + assertFalse(call.isDone(), "Call returned before waiting for accepted work"); + Thread.yield(); + } + fail("Caller did not begin waiting"); + } + + T result() throws Exception { + return call.get(WAIT.toMillis(), TimeUnit.MILLISECONDS); + } + + @Override + public void close() throws InterruptedException { + thread.interrupt(); + thread.join(WAIT.toMillis()); + assertFalse(thread.isAlive(), "Caller thread did not terminate"); + } } // ── S3 HTTP-level tests ─────────────────────────────────────────── @@ -307,11 +1061,15 @@ void requestUploadUrlDefaultsNullHeadersToEmptyMap() throws Exception { @Test void retryOnServerError() throws Exception { - var config = BraintrustConfig.builder().apiKey("test-api-key").apiUrl(baseUrl).build(); + var config = + BraintrustConfig.builder() + .apiKey("test-api-key") + .apiUrl(baseUrl) + .attachmentUploaderMaxRetries(2) + .attachmentUploaderInitialRetryDelay(Duration.ofMillis(100)) + .build(); var apiClient = BraintrustOpenApiClient.of(config); - var retryUploader = - new AttachmentUploader.S3AttachmentUploader( - apiClient, Duration.ofSeconds(30), 2, Duration.ofMillis(100)); + var retryUploader = new AttachmentUploader.S3AttachmentUploader(apiClient, config); // First two requests fail with 500, third succeeds stubFor(