fix: stop extension goroutine/memory leak and post-runtimeDone hang (#48)

flusher/flusher.go

@@ -1,10 +1,15 @@
 package flusher
 
 import (
+	"bytes"
+	"compress/gzip"
 	"context"
+	"encoding/json"
 	"fmt"
+	"io"
 	"log"
 	"os"
+	"strconv"
 	"sync"
 	"time"
 
@@ -29,15 +34,51 @@ var (
 	batchSize     = 1000
 	flushInterval = 1 * time.Second
 	logger        *zap.Logger
+
+	// maxBufferedEvents caps how many events may sit in the in-memory buffer.
+	// When ingestion fails, unsent events are requeued for a later attempt;
+	// without a cap, a sustained ingest outage on a high-volume function grows the
+	// buffer without bound and the extension leaks memory until it hits the Lambda
+	// memory ceiling (see
+	// https://github.com/axiomhq/axiom-lambda-extension/issues/48). Beyond the cap
+	// the oldest events are dropped. The default is ~10 flush batches, sized so the
+	// buffer stays small relative to the smallest (128MB) memory configurations.
+	// Override with AXIOM_MAX_BUFFERED_EVENTS.
+	maxBufferedEvents = 10_000
 )
 
 func init() {
 	logger, _ = zap.NewProduction()
+
+	if v := os.Getenv("AXIOM_MAX_BUFFERED_EVENTS"); v != "" {
+		if n, err := strconv.Atoi(v); err == nil && n > 0 {
+			maxBufferedEvents = n
+		} else {
+			logger.Warn("invalid AXIOM_MAX_BUFFERED_EVENTS, using default",
+				zap.String("value", v), zap.Int("default", maxBufferedEvents))
+		}
+	}
+}
+
+// ingester is the subset of *axiom.Client the flusher depends on. Depending on an
+// interface rather than the concrete client lets the flush path be unit-tested and
+// makes explicit that a flush honours the caller's context for cancellation.
+//
+// We use Ingest (a plain io.Reader body) rather than IngestEvents on purpose.
+// IngestEvents streams events through an io.Pipe fed by a background zstd encoder
+// goroutine; if the HTTP request is cancelled or stalls mid-flight, that goroutine
+// blocks forever in Encoder.Close -> PipeWriter.Write and leaks itself plus its
+// ~4MB compression buffer on every stalled flush. That is the memory leak in
+// https://github.com/axiomhq/axiom-lambda-extension/issues/48. Encoding the body
+// into a bytes buffer ourselves (see encodeBatch) has no background goroutine and
+// nothing that can be stranded.
+type ingester interface {
+	Ingest(ctx context.Context, id string, r io.Reader, typ axiom.ContentType, enc axiom.ContentEncoding, options ...ingest.Option) (*ingest.Status, error)
 }
 
 type Axiom struct {
-	client        *axiom.Client
-	retryClient   *axiom.Client
+	client        ingester
+	retryClient   ingester
 	events        []axiom.Event
 	eventsLock    sync.Mutex
 	lastFlushTime time.Time
@@ -95,24 +136,38 @@ func (f *Axiom) QueueEvents(events []axiom.Event) {
 	f.events = append(f.events, events...)
 }
 
-func (f *Axiom) Flush(opt RetryOpt) {
+// Flush sends the buffered events to Axiom. The provided context bounds the
+// ingest call: when it is cancelled (e.g. the per-invocation deadline is reached),
+// the in-flight request is aborted so the extension can hand control back to the
+// Lambda runtime instead of holding the sandbox open until the function times out
+// (see issue #48). On failure the batch is requeued for a later attempt, bounded
+// by maxBufferedEvents.
+func (f *Axiom) Flush(ctx context.Context, opt RetryOpt) {
 	f.eventsLock.Lock()
 	var batch []axiom.Event
 	// create a copy of the batch, clear the original
 	batch, f.events = f.events, []axiom.Event{}
+	f.lastFlushTime = time.Now()
 	f.eventsLock.Unlock()
 
-	f.lastFlushTime = time.Now()
 	if len(batch) == 0 {
 		return
 	}
 
+	body, err := encodeBatch(batch)
+	if err != nil {
+		// Encoding failure is not transient, but requeue (bounded) so a later
+		// flush can retry rather than silently dropping the batch.
+		logger.Error("Failed to encode events", zap.Error(err))
+		f.requeue(batch)
+		return
+	}
+
 	var res *ingest.Status
-	var err error
 	if opt == Retry {
-		res, err = f.retryClient.IngestEvents(context.Background(), axiomDataset, batch)
+		res, err = f.retryClient.Ingest(ctx, axiomDataset, body, axiom.NDJSON, axiom.Gzip)
 	} else {
-		res, err = f.client.IngestEvents(context.Background(), axiomDataset, batch)
+		res, err = f.client.Ingest(ctx, axiomDataset, body, axiom.NDJSON, axiom.Gzip)
 	}
 
 	if err != nil {
@@ -121,17 +176,63 @@ func (f *Axiom) Flush(opt RetryOpt) {
 		} else {
 			logger.Error("Failed to ingest events (will try again with next event)", zap.Error(err))
 		}
-		// allow this batch to be retried again, put them back
-		f.eventsLock.Lock()
-		defer f.eventsLock.Unlock()
-		f.events = append(batch, f.events...)
+		// Allow this batch to be retried again by putting it back in front of any
+		// events queued since, keeping the buffer bounded.
+		f.requeue(batch)
 
 		return
 	} else if res.Failed > 0 {
 		log.Printf("%d failures during ingesting, %s", res.Failed, res.Failures[0].Error)
 	}
 }
 
+// encodeBatch serialises events as gzip-compressed NDJSON into an in-memory
+// buffer. Building the body here (instead of via IngestEvents' streaming io.Pipe)
+// is what makes the flush leak-free: a bytes.Buffer never blocks, so the gzip
+// writer always closes cleanly and no encoder goroutine can be stranded when a
+// flush is cancelled (see issue #48 and the ingester doc above). The returned
+// *bytes.Reader lets net/http rewind the body for the shutdown retry path.
+func encodeBatch(batch []axiom.Event) (*bytes.Reader, error) {
+	var buf bytes.Buffer
+	gz := gzip.NewWriter(&buf)
+	enc := json.NewEncoder(gz)
+	for i := range batch {
+		if err := enc.Encode(batch[i]); err != nil {
+			_ = gz.Close()
+			return nil, err
+		}
+	}
+	if err := gz.Close(); err != nil {
+		return nil, err
+	}
+	return bytes.NewReader(buf.Bytes()), nil
+}
+
+// requeue puts a failed batch back at the front of the buffer, dropping the
+// oldest events when the buffer would exceed maxBufferedEvents. Dropping oldest
+// (rather than rejecting new) keeps the most recent logs, and copying into a
+// right-sized slice releases the dropped events' backing array to the GC so a
+// sustained outage cannot grow memory without bound (issue #48).
+func (f *Axiom) requeue(batch []axiom.Event) {
+	f.eventsLock.Lock()
+	combined := append(batch, f.events...)
+	dropped := 0
+	if len(combined) > maxBufferedEvents {
+		dropped = len(combined) - maxBufferedEvents
+		trimmed := make([]axiom.Event, maxBufferedEvents)
+		copy(trimmed, combined[dropped:])
+		combined = trimmed
+	}
+	f.events = combined
+	f.eventsLock.Unlock()
+
+	if dropped > 0 {
+		logger.Warn("event buffer full; dropped oldest events to bound memory (issue #48)",
+			zap.Int("dropped", dropped),
+			zap.Int("max_buffered_events", maxBufferedEvents))
+	}
+}
+
 // SafelyUseAxiomClient checks if axiom is empty, and if not, executes the given
 func SafelyUseAxiomClient(axiom *Axiom, action func(*Axiom)) {
 	if axiom != nil {

flusher/flusher_test.go

@@ -0,0 +1,148 @@
+package flusher
+
+import (
+	"context"
+	"errors"
+	"io"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/axiomhq/axiom-go/axiom"
+	"github.com/axiomhq/axiom-go/axiom/ingest"
+)
+
+// fakeIngester is a test double for the ingester interface.
+type fakeIngester struct {
+	mu    sync.Mutex
+	calls int
+	err   error
+	block bool // if true, block until the context is cancelled
+}
+
+func (f *fakeIngester) Ingest(ctx context.Context, _ string, r io.Reader, _ axiom.ContentType, _ axiom.ContentEncoding, _ ...ingest.Option) (*ingest.Status, error) {
+	f.mu.Lock()
+	f.calls++
+	f.mu.Unlock()
+
+	if f.block {
+		<-ctx.Done()
+		return nil, ctx.Err()
+	}
+	// Drain the body as a real client would.
+	_, _ = io.Copy(io.Discard, r)
+	if f.err != nil {
+		return nil, f.err
+	}
+	return &ingest.Status{}, nil
+}
+
+func (f *fakeIngester) callCount() int {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	return f.calls
+}
+
+func newTestAxiom(client ingester) *Axiom {
+	return &Axiom{
+		client:      client,
+		retryClient: client,
+		events:      make([]axiom.Event, 0),
+	}
+}
+
+// bufferLen returns the number of buffered events. Test helper.
+func (f *Axiom) bufferLen() int {
+	f.eventsLock.Lock()
+	defer f.eventsLock.Unlock()
+	return len(f.events)
+}
+
+func TestFlushClearsBufferOnSuccess(t *testing.T) {
+	fake := &fakeIngester{}
+	f := newTestAxiom(fake)
+	f.QueueEvents([]axiom.Event{{"a": 1}, {"b": 2}})
+
+	f.Flush(context.Background(), NoRetry)
+
+	if got := fake.callCount(); got != 1 {
+		t.Fatalf("expected 1 ingest call, got %d", got)
+	}
+	if n := f.bufferLen(); n != 0 {
+		t.Fatalf("expected empty buffer after successful flush, got %d", n)
+	}
+}
+
+func TestFlushRequeuesOnError(t *testing.T) {
+	fake := &fakeIngester{err: errors.New("boom")}
+	f := newTestAxiom(fake)
+	f.QueueEvents([]axiom.Event{{"a": 1}, {"b": 2}})
+
+	f.Flush(context.Background(), NoRetry)
+
+	if n := f.bufferLen(); n != 2 {
+		t.Fatalf("expected 2 events requeued after failed flush, got %d", n)
+	}
+}
+
+func TestFlushRequeueIsBounded(t *testing.T) {
+	prev := maxBufferedEvents
+	maxBufferedEvents = 3
+	defer func() { maxBufferedEvents = prev }()
+
+	fake := &fakeIngester{err: errors.New("boom")}
+	f := newTestAxiom(fake)
+	// Newer events are queued after the batch; the oldest must be dropped.
+	f.QueueEvents([]axiom.Event{{"n": 0}, {"n": 1}, {"n": 2}, {"n": 3}, {"n": 4}})
+
+	f.Flush(context.Background(), NoRetry)
+
+	if n := f.bufferLen(); n != 3 {
+		t.Fatalf("expected buffer capped at 3, got %d", n)
+	}
+	f.eventsLock.Lock()
+	defer f.eventsLock.Unlock()
+	if got := f.events[len(f.events)-1]["n"]; got != 4 {
+		t.Fatalf("expected newest event retained at tail, got %v", got)
+	}
+	if got := f.events[0]["n"]; got != 2 {
+		t.Fatalf("expected oldest retained to be n=2 after dropping 0 and 1, got %v", got)
+	}
+}
+
+func TestFlushRespectsContextCancellation(t *testing.T) {
+	fake := &fakeIngester{block: true}
+	f := newTestAxiom(fake)
+	f.QueueEvents([]axiom.Event{{"a": 1}})
+
+	ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
+	defer cancel()
+
+	done := make(chan struct{})
+	go func() {
+		f.Flush(ctx, NoRetry)
+		close(done)
+	}()
+
+	select {
+	case <-done:
+	case <-time.After(2 * time.Second):
+		t.Fatal("Flush did not return after context cancellation")
+	}
+
+	// Events should be requeued for a later attempt, not lost.
+	if n := f.bufferLen(); n != 1 {
+		t.Fatalf("expected event requeued after cancelled flush, got %d", n)
+	}
+}
+
+func TestFlushEmptyBufferIsNoop(t *testing.T) {
+	fake := &fakeIngester{}
+	f := newTestAxiom(fake)
+
+	f.Flush(context.Background(), NoRetry)
+
+	if got := fake.callCount(); got != 0 {
+		t.Fatalf("expected no ingest call for empty buffer, got %d", got)
+	}
+}