diff --git a/close_timeout_test.go b/close_timeout_test.go new file mode 100644 index 0000000..397ed40 --- /dev/null +++ b/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()) + } +} diff --git a/ingest_sync_test.go b/ingest_sync_test.go index ebe0e5c..b5fe3e7 100644 --- a/ingest_sync_test.go +++ b/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) diff --git a/quickwit.go b/quickwit.go index 3772182..6eacc38 100644 --- a/quickwit.go +++ b/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