feat: configurable SetCloseTimeout for the shutdown flush

close_timeout_test.go

@@ -0,0 +1,83 @@
+package quickwit_test
+
+import (
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/moonrhythm/quickwit"
+)
+
+// Core: the close-flush window rides out a transient failure during shutdown.
+// A fire-and-forget record is used so the FIRST flush happens during Close (a
+// tracked item would idle-flush earlier and be retried by the normal path, not
+// the close path). The server 5xxs the first attempts, then recovers within the
+// window, and the record must be delivered (not dropped).
+func TestClose_RidesOutTransientFailure(t *testing.T) {
+	var attempts, delivered atomic.Int64
+
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		io.Copy(io.Discard, r.Body)
+		if attempts.Add(1) <= 2 {
+			w.WriteHeader(http.StatusInternalServerError) // transient blip
+			return
+		}
+		delivered.Add(1)
+		io.WriteString(w, `{"num_ingested_docs":1,"num_rejected_docs":0}`)
+	}))
+	defer server.Close()
+
+	var discarded atomic.Int64
+	c := quickwit.NewClient(server.URL + "/api/v1/test")
+	c.SetConcurrent(1)
+	c.SetBatchSize(1000)            // never reached: flushed only at Close
+	c.SetMaxDelay(10 * time.Second) // never fires before Close
+	c.SetCloseTimeout(3 * time.Second)
+	c.OnDiscard(func(any) { discarded.Add(1) })
+
+	c.Ingest(map[string]any{"index": 0})
+	c.Close()
+
+	if delivered.Load() != 1 {
+		t.Errorf("delivered = %d, want 1 — Close should ride out the blip", delivered.Load())
+	}
+	if discarded.Load() != 0 {
+		t.Errorf("discarded = %d, want 0", discarded.Load())
+	}
+}
+
+// Core: Close gives up roughly at the close timeout when the server stays down,
+// discarding the buffered record rather than blocking forever or using the old
+// fixed ~0.8s budget. Fire-and-forget keeps the flush on the close path.
+func TestClose_GivesUpAtTimeout(t *testing.T) {
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		io.Copy(io.Discard, r.Body)
+		w.WriteHeader(http.StatusInternalServerError)
+	}))
+	defer server.Close()
+
+	var discarded atomic.Int64
+	c := quickwit.NewClient(server.URL + "/api/v1/test")
+	c.SetConcurrent(1)
+	c.SetBatchSize(1000)
+	c.SetMaxDelay(10 * time.Second)
+	c.SetCloseTimeout(400 * time.Millisecond)
+	c.OnDiscard(func(any) { discarded.Add(1) })
+
+	c.Ingest(map[string]any{"index": 0})
+
+	start := time.Now()
+	c.Close()
+	elapsed := time.Since(start)
+
+	// Bounded by the retry window (plus fast in-flight 500s); not unbounded.
+	if elapsed > 2*time.Second {
+		t.Errorf("Close took %v, want bounded near the 400ms close timeout", elapsed)
+	}
+	if discarded.Load() != 1 {
+		t.Errorf("OnDiscard fired %d times, want 1 (record dropped after the window)", discarded.Load())
+	}
+}

ingest_sync_test.go

@@ -281,12 +281,13 @@ func TestIngestBatch_ClosedWhileServerFailsSettlesClosed(t *testing.T) {
 
 	c := quickwit.NewClient(server.URL + "/api/v1/test")
 	c.SetConcurrent(1)
-	c.SetBatchSize(1000)            // never reached
-	c.SetMaxDelay(10 * time.Second) // never fires before Close
+	c.SetBatchSize(1000)                      // never reached
+	c.SetMaxDelay(10 * time.Second)           // never fires before Close
+	c.SetCloseTimeout(200 * time.Millisecond) // give up quickly for the test
 	r := c.IngestBatch(map[string]any{"index": 0})
 
-	// Close drains the buffer, fails the final flush 5×, then gives up and
-	// settles the remaining items.
+	// Close drains the buffer, retries the final flush until the close timeout,
+	// then gives up and settles the remaining items.
 	c.Close()
 
 	ctx, cancel := context.WithTimeout(context.Background(), time.Second)

quickwit.go

@@ -26,6 +26,7 @@ const (
 	ReduceBatchSizeToRatio = 0.9 // reduce 10% of the batch size
 	ReduceBatchSizeMin     = 0.1 // do not reduce below 10% of the default batch size
 	ResetBatchSizeAfter    = 10 * time.Minute
+	CloseTimeout           = 5 * time.Second // how long Close retries the final flush before discarding
 
 	// ingestResponseMaxBytes caps how much of the ingest response body is read
 	// before parsing, so a pathological body cannot exhaust memory. The default
@@ -226,6 +227,7 @@ type Client struct {
 	maxDelay            time.Duration
 	ingestBufferSize    int
 	ingestTimeout       time.Duration
+	closeTimeout        time.Duration
 	discard             bool
 	concurrent          int
 	ingestBuffer        chan ingestItem
@@ -275,6 +277,21 @@ func (c *Client) SetIngestTimeout(timeout time.Duration) {
 	c.ingestTimeout = timeout
 }
 
+// SetCloseTimeout bounds the retry window Close spends re-attempting the final
+// flush of buffered records before giving up and discarding them. A larger value
+// rides out a transient outage (e.g. a rolling restart); a permanent rejection
+// still exits fast. Records still undelivered when it elapses are dropped
+// (OnDiscard for fire-and-forget, settled ReasonClosed for tracked).
+//
+// It bounds the retry window, not total Close time: each attempt makes a real
+// request bounded by SetIngestTimeout, so an in-flight flush can overrun the
+// window by up to that, and a record already mid-retry when Close is called can
+// take up to roughly twice the window. Set before the first Ingest; defaults to
+// CloseTimeout.
+func (c *Client) SetCloseTimeout(timeout time.Duration) {
+	c.closeTimeout = timeout
+}
+
 func (c *Client) SetDiscard(discard bool) {
 	c.discard = discard
 }
@@ -397,6 +414,13 @@ func (c *Client) getMaxDelay() time.Duration {
 	return c.maxDelay
 }
 
+func (c *Client) getCloseTimeout() time.Duration {
+	if c.closeTimeout <= 0 {
+		return CloseTimeout
+	}
+	return c.closeTimeout
+}
+
 func (c *Client) getBatchSize() int {
 	if c.batchSize <= 0 {
 		return IngestBatchSize
@@ -650,6 +674,10 @@ func (c *Client) discardItem(it ingestItem, err error) {
 	c.invokeOnDiscard(it.data)
 }
 
+// Close stops the workers after flushing what is buffered. It blocks until the
+// final flush succeeds or the close-timeout retry window (see SetCloseTimeout)
+// elapses, after which any still-undelivered records are discarded. Close is
+// idempotent and safe to call before the first Ingest.
 func (c *Client) Close() {
 	c.onceSetup.Do(c.setup)
 	c.onceClose.Do(func() {
@@ -958,30 +986,37 @@ func (c *Client) loop() {
 
 	closing:
 
-		// For closing case, use limited retries
-		maxRetries := 5
+		// For the closing case, retry until the close-timeout deadline. A
+		// permanent rejection still exits immediately (flush reports it handled);
+		// only transient failures consume the window. Retry-After is intentionally
+		// not honored here so a server-requested delay cannot stretch shutdown.
+		deadline := time.Now().Add(c.getCloseTimeout())
 		backoff := 100 * time.Millisecond
+		attempt := 0
 
-		for i := 0; i < maxRetries; i++ {
+		for {
 			if flushOversize() {
 				return true
 			}
 
-			// Don't sleep on the last attempt
-			if i < maxRetries-1 {
-				slog.Info("quickwit: flush failed while closing, retrying", "attempt", i+1, "maxRetries", maxRetries)
-				// Jittered backoff; Retry-After is intentionally not honored here
-				// so a server-requested delay cannot stretch shutdown.
-				time.Sleep(jitterBackoff(backoff))
-				// Exponential backoff with a cap
-				backoff = time.Duration(float64(backoff) * 1.5)
-				if backoff > time.Second {
-					backoff = time.Second
-				}
+			remaining := time.Until(deadline)
+			if remaining <= 0 {
+				return false
 			}
-		}
 
-		return false
+			attempt++
+			slog.Info("quickwit: flush failed while closing, retrying", "attempt", attempt)
+			sleep := jitterBackoff(backoff)
+			if sleep > remaining {
+				sleep = remaining // never sleep past the deadline
+			}
+			time.Sleep(sleep)
+			// Exponential backoff with a cap
+			backoff = time.Duration(float64(backoff) * 1.5)
+			if backoff > time.Second {
+				backoff = time.Second
+			}
+		}
 	}
 
 	// finalFlush drains the buffer on shutdown, settling/discarding anything the