Allow forwarder to be stopped cleanly

dogstatsd-http/forwarder/src/main/java/com/datadoghq/dogstatsd/http/forwarder/BoundedQueue.java

@@ -62,6 +62,7 @@ public int compareTo(Key o) {
     Lock lock = new ReentrantLock();
     Condition notEmpty = lock.newCondition();
     Condition notFull = lock.newCondition();
+    boolean closed;
 
     BoundedQueue(long maxBytes, long maxTries, WhenFull whenFull, Telemetry telemetry) {
         this(maxBytes, maxTries, whenFull, telemetry, System::nanoTime);
@@ -111,6 +112,11 @@ private void put(Key key, Payload item, WhenFull whenFull) throws InterruptedExc
         lock.lock();
         try {
             if (key == null) {
+                // Queue is closed for new items, but retries are allowed to enter.
+                if (closed) {
+                    throw new IllegalStateException("closed");
+                }
+
                 key = newKey();
             }
             ensureSpace(item.bytes.length, whenFull);
@@ -141,6 +147,9 @@ private void ensureSpace(int length, WhenFull whenFull) throws InterruptedExcept
                     break;
                 case BLOCK:
                     notFull.await();
+                    if (closed) {
+                        throw new IllegalStateException("closed");
+                    }
                     break;
             }
         }
@@ -149,7 +158,10 @@ private void ensureSpace(int length, WhenFull whenFull) throws InterruptedExcept
     Map.Entry<Key, Payload> next() throws InterruptedException {
         lock.lock();
         try {
-            while (items.size() == 0) {
+            while (empty()) {
+                if (closed) {
+                    return null;
+                }
                 notEmpty.await();
             }
             Map.Entry<Key, Payload> item = items.pollFirstEntry();
@@ -161,6 +173,17 @@ Map.Entry<Key, Payload> next() throws InterruptedException {
         }
     }
 
+    void close() {
+        lock.lock();
+        try {
+            closed = true;
+            notEmpty.signalAll();
+            notFull.signalAll();
+        } finally {
+            lock.unlock();
+        }
+    }
+
     void snapshot(long now, Telemetry.Snapshot s) {
         lock.lock();
         try {
@@ -179,4 +202,8 @@ void snapshot(long now, Telemetry.Snapshot s) {
             lock.unlock();
         }
     }
+
+    boolean empty() {
+        return items.size() == 0;
+    }
 }

dogstatsd-http/forwarder/src/main/java/com/datadoghq/dogstatsd/http/forwarder/Forwarder.java

@@ -79,9 +79,13 @@ public Telemetry.Snapshot snapshot() {
     /** Runs the forwarding loop, delivering queued payloads until the thread is interrupted. */
     @Override
     public void run() {
-        while (true) {
+        while (!interrupted()) {
             try {
-                runOnce(queue.next());
+                Map.Entry<BoundedQueue.Key, Payload> item = queue.next();
+                if (item == null) {
+                    return;
+                }
+                runOnce(item);
             } catch (InterruptedException e) {
                 return;
             } catch (Exception t) {
@@ -100,6 +104,7 @@ public void run() {
      * @param payload the raw bytes to deliver
      * @throws InterruptedException if the calling thread is interrupted while waiting for space
      *     ({@link WhenFull#BLOCK} mode only)
+     * @throws IllegalStateException if the forwarder has been closed via {@link #close(Duration)}
      */
     public void send(URI url, byte[] payload) throws InterruptedException {
         Objects.requireNonNull(url, "url");
@@ -139,6 +144,10 @@ void runOnce(Map.Entry<BoundedQueue.Key, Payload> item) throws InterruptedExcept
         } catch (IOException ex) {
             logger.log(Level.WARNING, "error sending request: {0}", ex.toString());
             handleTransportError(item);
+        } catch (InterruptedException ex) {
+            // Wouldn't be retried, but will show up as a leftover in a telemetry snapshot.
+            queue.requeue(item);
+            throw ex;
         }
 
         backoff();
@@ -229,4 +238,34 @@ private static void validateHeaderValue(String value) {
             throw new IllegalArgumentException("invalid character");
         }
     }
+
+    /**
+     * Closes the forwarder: stops accepting new payloads and drains the remaining backlog.
+     *
+     * <p>Already-queued payloads keep being delivered until either the queue drains or {@code
+     * timeout} elapses, whichever comes first. If the timeout elapses first the forwarding thread
+     * is interrupted, abandoning any unsent payloads.
+     *
+     * @param timeout maximum time to wait for the backlog to drain. {@code null} means wait
+     *     forever.
+     * @return {@code true} if the queue drained cleanly with no unsent payloads remaining; {@code
+     *     false} if the timeout elapsed with data still queued
+     * @throws InterruptedException if the calling thread is interrupted while waiting
+     */
+    public boolean close(Duration timeout) throws InterruptedException {
+        queue.close();
+        if (timeout == null) {
+            join(0);
+        } else {
+            long timeoutMs = timeout.toMillis();
+            if (timeoutMs > 0) {
+                join(timeoutMs);
+            }
+        }
+        if (isAlive()) {
+            interrupt();
+            join();
+        }
+        return queue.empty();
+    }
 }

dogstatsd-http/forwarder/src/test/java/com/datadoghq/dogstatsd/http/forwarder/BoundedQueueTest.java

@@ -148,6 +148,49 @@ public void blockUnblocksWhenSpaceFreed() throws InterruptedException {
         assertEquals(0, t.snapshot(q).droppedPayloads);
     }
 
+    @Test(timeout = 5000)
+    public void closeUnblocksProducerWithException() throws InterruptedException {
+        BoundedQueue q = new BoundedQueue(5, 1, WhenFull.BLOCK, new Telemetry());
+        q.add(payload(5)); // queue full
+
+        Throwable[] thrown = new Throwable[1];
+        Thread producer =
+                new Thread(
+                        () -> {
+                            try {
+                                q.add(payload(5));
+                            } catch (Throwable t) {
+                                thrown[0] = t;
+                            }
+                        });
+        producer.start();
+
+        // Give the producer time to reach await().
+        while (!(producer.getState() == Thread.State.WAITING
+                || producer.getState() == Thread.State.TIMED_WAITING)) {
+            Thread.sleep(50);
+        }
+
+        q.close();
+        producer.join(2000);
+        assertFalse(producer.isAlive());
+        assertTrue(thrown[0] instanceof IllegalStateException);
+    }
+
+    @Test(expected = IllegalStateException.class)
+    public void addAfterCloseThrows() throws InterruptedException {
+        BoundedQueue q = new BoundedQueue(10, 1, WhenFull.DROP, new Telemetry());
+        q.close();
+        q.add(payload(3));
+    }
+
+    @Test
+    public void nextReturnsNullWhenClosedAndEmpty() throws InterruptedException {
+        BoundedQueue q = new BoundedQueue(10, 1, WhenFull.DROP, new Telemetry());
+        q.close();
+        assertNull(q.next());
+    }
+
     @Test(expected = IllegalArgumentException.class)
     public void addThrowsForOversizedItem() throws InterruptedException {
         BoundedQueue q = new BoundedQueue(4, 1, WhenFull.DROP, new Telemetry());

dogstatsd-http/forwarder/src/test/java/com/datadoghq/dogstatsd/http/forwarder/ForwarderTest.java

@@ -8,12 +8,16 @@
 package com.datadoghq.dogstatsd.http.forwarder;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertThrows;
+import static org.junit.Assert.assertTrue;
 
 import java.net.URI;
 import java.time.Duration;
 import java.util.Map;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.atomic.AtomicInteger;
 import org.junit.Test;
 
 public class ForwarderTest {
@@ -109,4 +113,66 @@ public void handle500() throws Exception {
         assertEquals(1, cc.payloads);
         assertEquals(7, cc.bytes);
     }
+
+    /** With an empty queue, close() unblocks the loop's next() and drains cleanly. */
+    @Test(timeout = 5000)
+    public void closeDrainsEmptyQueueReturnsTrue() throws InterruptedException {
+        Forwarder f = newForwarder(100, WhenFull.DROP);
+        f.start();
+        assertTrue(f.close(Duration.ofSeconds(5)));
+        assertFalse(f.isAlive());
+    }
+
+    /** Payloads queued before close() are all delivered before the thread exits. */
+    @Test(timeout = 5000)
+    public void closeDrainsPendingItemsReturnsTrue() throws InterruptedException {
+        AtomicInteger processed = new AtomicInteger();
+        Forwarder f =
+                new Forwarder(100, 1, WhenFull.DROP, Duration.ofSeconds(1), Duration.ofSeconds(1)) {
+                    @Override
+                    void runOnce(Map.Entry<BoundedQueue.Key, Payload> item) {
+                        processed.incrementAndGet();
+                    }
+                };
+        f.send(URL, new byte[3]);
+        f.send(URL, new byte[3]);
+        f.send(URL, new byte[3]);
+        f.start();
+        assertTrue(f.close(Duration.ofSeconds(5)));
+        assertFalse(f.isAlive());
+        assertEquals(3, processed.get());
+    }
+
+    /** After close(), send() propagates the closed-queue IllegalStateException. */
+    @Test
+    public void sendAfterCloseThrows() throws InterruptedException {
+        Forwarder f = newForwarder(100, WhenFull.DROP);
+        assertTrue(f.close(Duration.ofSeconds(1)));
+        assertThrows(IllegalStateException.class, () -> f.send(URL, new byte[3]));
+    }
+
+    /** If the backlog can't drain in time, close() interrupts the thread and returns false. */
+    @Test(timeout = 5000)
+    public void closeTimesOutReturnsFalse() throws InterruptedException {
+        CountDownLatch entered = new CountDownLatch(1);
+        Forwarder f =
+                new Forwarder(100, 1, WhenFull.DROP, Duration.ofSeconds(1), Duration.ofSeconds(1)) {
+                    @Override
+                    void runOnce(Map.Entry<BoundedQueue.Key, Payload> item)
+                            throws InterruptedException {
+                        try {
+                            entered.countDown();
+                            Thread.sleep(Long.MAX_VALUE);
+                        } catch (InterruptedException ex) {
+                            queue.requeue(item);
+                            throw ex;
+                        }
+                    }
+                };
+        f.send(URL, new byte[3]);
+        f.start();
+        entered.await();
+        assertFalse(f.close(Duration.ofMillis(200)));
+        assertFalse(f.isAlive());
+    }
 }