From 8d4e4866054a61d2961db911b3e93c2e00531182 Mon Sep 17 00:00:00 2001 From: Sean Linsley Date: Mon, 13 Jul 2026 13:44:40 -0400 Subject: [PATCH 1/3] Retry failed snapshot submission using internal queue --- output/compact.go | 5 +- output/full.go | 4 +- output/upload.go | 99 ++++++++-------- output/upload_http_legacy.go | 125 ++------------------ runner/run.go | 3 + state/memory_limit.go | 36 ++++++ state/queue.go | 214 +++++++++++++++++++++++++++++++++++ state/queue_test.go | 169 +++++++++++++++++++++++++++ state/state.go | 6 +- util/reconnecting_socket.go | 19 +++- 10 files changed, 506 insertions(+), 174 deletions(-) create mode 100644 state/memory_limit.go create mode 100644 state/queue.go create mode 100644 state/queue_test.go diff --git a/output/compact.go b/output/compact.go index fec19acbd..877cd5f4f 100644 --- a/output/compact.go +++ b/output/compact.go @@ -39,9 +39,8 @@ func uploadAndSubmitCompactSnapshot(ctx context.Context, s *pganalyze_collector. return nil } - server.CompactSnapshotUpload <- s - - return nil + kind := kindFromCompactSnapshot(s) + return server.SnapshotQueue.Push(kind, s) } func kindFromCompactSnapshot(s *pganalyze_collector.CompactSnapshot) string { diff --git a/output/full.go b/output/full.go index 18a124dc6..08d0e07ad 100644 --- a/output/full.go +++ b/output/full.go @@ -65,9 +65,7 @@ func submitFull(ctx context.Context, s *snapshot.FullSnapshot, server *state.Ser return nil } - server.FullSnapshotUpload <- s - - return nil + return server.SnapshotQueue.Push("full", s) } func verifyIntegrity(s *snapshot.FullSnapshot) error { diff --git a/output/upload.go b/output/upload.go index 0c7b0c45b..5875bc2ee 100644 --- a/output/upload.go +++ b/output/upload.go @@ -1,8 +1,6 @@ package output import ( - "bytes" - "compress/zlib" "context" "errors" "fmt" @@ -11,7 +9,6 @@ import ( "github.com/pganalyze/collector/state" "github.com/pganalyze/collector/util" - "google.golang.org/protobuf/proto" ) func SetupSnapshotUploadForAllServers(ctx context.Context, servers []*state.Server, opts state.CollectionOpts, logger *util.Logger) { @@ -19,47 +16,51 @@ func SetupSnapshotUploadForAllServers(ctx context.Context, servers []*state.Serv return } for _, server := range servers { - go snapshotUploadForServer(ctx, server, logger.WithPrefixAndRememberErrors(server.Config.SectionName), opts) + prefixedLogger := logger.WithPrefixAndRememberErrors(server.Config.SectionName) + server.SnapshotQueue.Logger = prefixedLogger + go snapshotUploadForServer(ctx, server, prefixedLogger, opts) } } func snapshotUploadForServer(ctx context.Context, server *state.Server, logger *util.Logger, opts state.CollectionOpts) { var compactLogTime time.Time - compactLogStats := make(map[string]uint8) + var compactLogStats = make(map[string]uint8) + var failed bool + var delay time.Duration + for { + if failed { + delay = min(5*delay, 10*time.Second) // Increasing backoff delay in case of failure + } else { + delay = 10 * time.Millisecond // Small delay to avoid high CPU usage in loop + } select { case <-ctx.Done(): return - case s := <-server.FullSnapshotUpload: - data, err := proto.Marshal(s) - if err != nil { - logger.PrintError("Error marshaling protocol buffers") - continue - } + case <-time.After(delay): + } - err = uploadViaWebsocketOrHttp(ctx, server, logger, opts, data, s.SnapshotUuid, s.CollectedAt.AsTime(), false) - if err != nil { - logger.PrintError("Error uploading snapshot: %s", err) - } else if !opts.TestRun { - logger.PrintInfo("Submitted full snapshot successfully") - } - case s := <-server.CompactSnapshotUpload: - data, err := proto.Marshal(s) - if err != nil { - logger.PrintError("Error marshaling protocol buffers") - continue - } + tx, err := server.SnapshotQueue.Pop(ctx) + if err != nil { + continue + } - err = uploadViaWebsocketOrHttp(ctx, server, logger, opts, data, s.SnapshotUuid, s.CollectedAt.AsTime(), false) - if err != nil { - logger.PrintError("Error uploading snapshot: %s", err) - continue + err = uploadViaWebsocketOrHttp(ctx, server, logger, opts, tx.Snapshot) + if err != nil { + logger.PrintError("Error uploading %s snapshot: %s", tx.Kind, err) + tx.Rollback() + failed = true + } else { + tx.Commit() + failed = false + if !opts.TestRun { + logger.PrintInfo("Submitted %s snapshot successfully", tx.Kind) } - if opts.TestRun { + if tx.Kind == "full" { continue } - - kind := kindFromCompactSnapshot(s) + // Compact snapshot: log stats periodically + kind := tx.Kind logger.PrintVerbose("Submitted compact %s snapshot successfully", kind) compactLogStats[kind] = compactLogStats[kind] + 1 if compactLogTime.IsZero() { @@ -92,23 +93,31 @@ func summarizeCounts(counts map[string]uint8) string { return details } -func uploadViaWebsocketOrHttp(ctx context.Context, server *state.Server, logger *util.Logger, opts state.CollectionOpts, data []byte, snapshotUUID string, collectedAt time.Time, compactSnapshot bool) error { - var compressedData bytes.Buffer - w := zlib.NewWriter(&compressedData) - w.Write(data) - w.Close() - +func uploadViaWebsocketOrHttp(ctx context.Context, server *state.Server, logger *util.Logger, opts state.CollectionOpts, data []byte) error { if server.WebSocket.Connected() { logger.PrintVerbose("Uploading snapshot to websocket") - server.WebSocket.Write <- compressedData.Bytes() - } else if server.Config.APIRequireWebsocket { - return errors.New("Error uploading snapshot: WebSocket not connected") - } else { - s3Location, err := uploadSnapshot(ctx, server.Config.HTTPClientWithRetry, server.Grant.Load(), logger, compressedData.Bytes(), snapshotUUID) - if err != nil { - return err + result := make(chan error, 1) + select { + case server.WebSocket.Write <- util.WriteRequest{Data: data, Result: result}: + select { + case err := <-result: + if err != nil { + return fmt.Errorf("WebSocket write failed: %w", err) + } + return nil + case <-time.After(5 * time.Second): + logger.PrintWarning("WebSocket write timed out, falling back to HTTP") + case <-ctx.Done(): + return ctx.Err() + } + case <-time.After(5 * time.Second): + logger.PrintWarning("WebSocket write timed out, falling back to HTTP") + case <-ctx.Done(): + return ctx.Err() } - submitSnapshot(ctx, server, opts, logger, s3Location, collectedAt, compactSnapshot) } - return nil + if server.Config.APIRequireWebsocket { + return errors.New("Error uploading snapshot: WebSocket not connected") + } + return uploadSnapshot(ctx, server.Config.HTTPClient, server.Grant.Load(), logger, data) } diff --git a/output/upload_http_legacy.go b/output/upload_http_legacy.go index d7630e302..dafdc62d9 100644 --- a/output/upload_http_legacy.go +++ b/output/upload_http_legacy.go @@ -3,59 +3,27 @@ package output import ( "bytes" "context" - "encoding/json" - "encoding/xml" "fmt" "io" - "mime" "mime/multipart" "net/http" - "net/url" - "os" - "path/filepath" - "strings" "time" - "github.com/pganalyze/collector/config" "github.com/pganalyze/collector/state" "github.com/pganalyze/collector/util" ) -func uploadSnapshot(ctx context.Context, httpClient *http.Client, grant *state.Grant, logger *util.Logger, data []byte, filename string) (string, error) { - var err error - +func uploadSnapshot(ctx context.Context, httpClient *http.Client, grant *state.Grant, logger *util.Logger, data []byte) error { if !grant.ValidForS3Until.After(time.Now()) { - return "", fmt.Errorf("Error - can't upload without valid S3 grant") - } - - if grant.S3URL == "" && grant.LocalDir != "" { - location := grant.LocalDir + filename - err = os.MkdirAll(filepath.Dir(location), 0755) - if err != nil { - logger.PrintError("Error creating target directory: %s", err) - return "", err - } - - err = os.WriteFile(location, data, 0644) - if err != nil { - logger.PrintError("Error writing local file: %s", err) - return "", err - } - return location, nil + return fmt.Errorf("Error - can't upload without valid S3 grant") } logger.PrintVerbose("Successfully prepared S3 request - size of request body: %.4f MB", float64(len(data))/1024.0/1024.0) - return uploadToS3(ctx, httpClient, grant.S3URL, grant.S3Fields, data, filename) + return uploadToS3(ctx, httpClient, grant.S3URL, grant.S3Fields, data) } -type s3UploadResponse struct { - Location string - Bucket string - Key string -} - -func uploadToS3(ctx context.Context, httpClient *http.Client, S3URL string, S3Fields map[string]string, data []byte, filename string) (string, error) { +func uploadToS3(ctx context.Context, httpClient *http.Client, S3URL string, S3Fields map[string]string, data []byte) error { var err error var formBytes bytes.Buffer @@ -64,75 +32,28 @@ func uploadToS3(ctx context.Context, httpClient *http.Client, S3URL string, S3Fi for key, val := range S3Fields { err = writer.WriteField(key, val) if err != nil { - return "", err + return err } } - part, _ := writer.CreateFormFile("file", filename) + part, _ := writer.CreateFormFile("file", "snapshot") _, err = part.Write(data) if err != nil { - return "", err + return err } writer.Close() req, err := http.NewRequestWithContext(ctx, "POST", S3URL, &formBytes) if err != nil { - return "", err + return err } req.Header.Set("Content-Type", writer.FormDataContentType()) resp, err := httpClient.Do(req) - if err != nil { - return "", err - } - defer resp.Body.Close() - - body, err := io.ReadAll(resp.Body) - if err != nil { - return "", err - } - - if resp.StatusCode != http.StatusCreated { - return "", fmt.Errorf("Bad S3 upload return code %s (expected 201 Created), body: %s", resp.Status, body) - } - - var s3Resp s3UploadResponse - err = xml.Unmarshal(body, &s3Resp) - if err != nil { - return "", err - } - - return s3Resp.Key, nil -} - -func submitSnapshot(ctx context.Context, server *state.Server, opts state.CollectionOpts, logger *util.Logger, s3Location string, collectedAt time.Time, compact bool) error { - requestURL := server.Config.APIBaseURL + "/v2/snapshots" - - if opts.TestRun { - requestURL = server.Config.APIBaseURL + "/v2/snapshots/test" - } else if compact { - requestURL = server.Config.APIBaseURL + "/v2/snapshots/compact" - } - - data := url.Values{ - "s3_location": {s3Location}, - "collected_at": {fmt.Sprintf("%d", collectedAt.Unix())}, - } - - req, err := http.NewRequestWithContext(ctx, "POST", requestURL, strings.NewReader(data.Encode())) if err != nil { return err } - - req.Header = config.APIHeaders(server.Config, opts.TestRun, opts.StartedAt) - req.Header.Set("Content-Type", "application/x-www-form-urlencoded") - req.Header.Add("Accept", "application/json,text/plain") - - resp, err := server.Config.HTTPClientWithRetry.Do(req) - if err != nil { - return util.CleanHTTPError(err) - } defer resp.Body.Close() body, err := io.ReadAll(resp.Body) @@ -140,34 +61,8 @@ func submitSnapshot(ctx context.Context, server *state.Server, opts state.Collec return err } - if resp.StatusCode != http.StatusOK { - return fmt.Errorf("Error when submitting: %s\n", body) - } - - if opts.TestRun { - contentType, _, err := mime.ParseMediaType(resp.Header.Get("Content-Type")) - if err != nil { - return fmt.Errorf("Error decoding response: %s\n", err) - } - - var msg string - - if contentType == "application/json" { - var jsonBody struct { - Message string `json:"message"` - } - err = json.Unmarshal(body, &jsonBody) - if err != nil { - return fmt.Errorf("Error decoding response: %s\n", err) - } - msg = jsonBody.Message - } else { - msg = string(body) - } - - if len(msg) > 0 { - logger.PrintInfo(" %s", msg) - } + if resp.StatusCode != http.StatusCreated { + return fmt.Errorf("Bad S3 upload return code %s (expected 201 Created), body: %s", resp.Status, body) } return nil diff --git a/runner/run.go b/runner/run.go index 58579ca5f..c20ab0018 100644 --- a/runner/run.go +++ b/runner/run.go @@ -91,6 +91,9 @@ func Run(ctx context.Context, wg *sync.WaitGroup, opts state.CollectionOpts, log } shutdown = func() { + for _, server := range servers { + server.SnapshotQueue.Close() + } for _, cfg := range conf.Servers { if cfg.OTelTracingProviderShutdownFunc == nil { continue diff --git a/state/memory_limit.go b/state/memory_limit.go new file mode 100644 index 000000000..e53ab6c22 --- /dev/null +++ b/state/memory_limit.go @@ -0,0 +1,36 @@ +package state + +import ( + "sync/atomic" +) + +// MemoryLimit tracks a shared byte counter against a configurable cap. +// Callers use Add/Remove to adjust the counter and Size to inspect it. +// OverLimit reports whether the current usage has exceeded the cap. +type MemoryLimit struct { + bytes atomic.Int64 + limit int64 +} + +// Global memory limit for all snapshot queues +var QueueMemory = NewMemoryLimit(200 * 1024 * 1024) + +func NewMemoryLimit(cap int64) *MemoryLimit { + return &MemoryLimit{limit: cap} +} + +func (m *MemoryLimit) Add(n int64) int64 { + return m.bytes.Add(n) +} + +func (m *MemoryLimit) Remove(n int64) int64 { + return m.bytes.Add(-n) +} + +func (m *MemoryLimit) Size() int64 { + return m.bytes.Load() +} + +func (m *MemoryLimit) OverLimit() bool { + return m.bytes.Load() > m.limit +} diff --git a/state/queue.go b/state/queue.go new file mode 100644 index 000000000..07d1e7248 --- /dev/null +++ b/state/queue.go @@ -0,0 +1,214 @@ +package state + +import ( + "bytes" + "compress/zlib" + "context" + "errors" + "sync" + + "github.com/pganalyze/collector/util" + "google.golang.org/protobuf/proto" +) + +const DefaultCapacity = 500 + +// Queue of snapshots ready for submission to pganalyze +// - Thread safe: Safe for concurrent use by multiple readers and writers. +// - Limited capacity: If full, Push drops the oldest item to make room. +// - Transactional: Pop locks the head item via a generation ID which is removed on Commit, and re-released on Rollback. +// - Generation tracking: Increments a counter on every Push to uniquely identify each item. +// If a Push evicts an in-flight item, its generation changes, safely ignoring late Commits or Rollbacks. +// - Memory-bounded: When the global QueueMemory limit is exceeded, Push drops the oldest items +// to stay within the budget. All queues share this single global limit. +type Queue struct { + mu sync.Mutex + cond *sync.Cond + data []QueueItem + capacity int + head int + tail int + size int + sizeBytes int64 + activeGen uint64 + closed bool + nextGenID uint64 + Logger *util.Logger + DropCallback func(kind string, sizeBytes int64) +} + +func NewQueue(logger *util.Logger) *Queue { + q := &Queue{ + data: make([]QueueItem, DefaultCapacity), + capacity: DefaultCapacity, + Logger: logger, + } + q.cond = sync.NewCond(&q.mu) + return q +} + +// Push adds an item to the tail. If full, it drops the oldest item to make room. +// When the global queue memory limit is exceeded, Push drops the oldest items +// (starting with the current queue's oldest) to stay within the budget. +// All queues share the single global memory limit defined by QueueMemory. +func (q *Queue) Push(kind string, snapshot proto.Message) (err error) { + data, err := proto.Marshal(snapshot) + if err != nil { + return + } + var buf bytes.Buffer + w := zlib.NewWriter(&buf) + w.Write(data) + w.Close() + return q.PushBytes(kind, buf.Bytes()) +} + +// Factored out from Push so it can be called by tests without constructing real snapshots +func (q *Queue) PushBytes(kind string, bytes []byte) (err error) { + sizeBytes := int64(len(bytes)) + q.mu.Lock() + defer q.mu.Unlock() + if q.closed { + return + } + q.makeSpace(sizeBytes) + q.data[q.tail] = QueueItem{ + Kind: kind, + Snapshot: bytes, + SizeBytes: sizeBytes, + Generation: q.nextGenID, + } + q.tail = (q.tail + 1) % q.capacity + q.size++ + q.sizeBytes += sizeBytes + q.cond.Signal() + return +} + +// Pop blocks until an item is ready or the queue closes. +// On success, it locks the head item and returns a Transaction handle. +// +// The caller is woken by Push (Signal), Commit/Rollback (Broadcast), or Close (Broadcast). +func (q *Queue) Pop(ctx context.Context) (*Transaction, error) { + q.mu.Lock() + defer q.mu.Unlock() + for (q.size == 0 || q.activeGen != 0) && !q.closed && ctx.Err() == nil { + q.cond.Wait() + } + if ctx.Err() != nil { + return nil, ctx.Err() + } + if q.closed { + return nil, errors.New("queue closed") + } + item := q.data[q.head] + q.activeGen = item.Generation + return &Transaction{ + Kind: item.Kind, + Snapshot: item.Snapshot, + SizeBytes: item.SizeBytes, + generation: item.Generation, + q: q, + }, nil +} + +// Close terminates the queue, unblocks waiting readers, and rejects new writes +func (q *Queue) Close() { + q.mu.Lock() + if !q.closed { + q.closed = true + q.cond.Broadcast() + } + q.mu.Unlock() +} + +type QueueItem struct { + Kind string + Snapshot []byte + SizeBytes int64 + Generation uint64 +} + +type Transaction struct { + Kind string + Snapshot []byte + SizeBytes int64 + generation uint64 + q *Queue +} + +func (t *Transaction) Commit() { + t.q.mu.Lock() + defer t.q.mu.Unlock() + // Validate that the transaction hasn't been evicted or superseded + if t.q.activeGen != t.generation || t.q.data[t.q.head].Generation != t.generation { + return + } + t.q.data[t.q.head] = QueueItem{} + t.q.head = (t.q.head + 1) % t.q.capacity + t.q.size-- + t.q.sizeBytes -= t.SizeBytes + QueueMemory.Remove(t.SizeBytes) + t.q.activeGen = 0 + t.q.cond.Broadcast() +} + +func (t *Transaction) Rollback() { + t.q.mu.Lock() + defer t.q.mu.Unlock() + // Validate that the transaction hasn't been evicted or superseded + if t.q.activeGen != t.generation || t.q.data[t.q.head].Generation != t.generation { + return + } + t.q.activeGen = 0 + t.q.cond.Broadcast() +} + +func (q *Queue) makeSpace(sizeBytes int64) { + QueueMemory.Add(sizeBytes) + // Evict items to satisfy the global memory limit + evictedCount := 0 + for evictedCount <= 100 { + if !QueueMemory.OverLimit() { + break + } + if q.size == 0 { + // This queue is empty; another server's queue is the problem + break + } + evicted := q.data[q.head] + if q.activeGen == evicted.Generation { + q.activeGen = 0 + } + q.data[q.head] = QueueItem{} + q.head = (q.head + 1) % q.capacity + q.size-- + q.sizeBytes -= evicted.SizeBytes + QueueMemory.Remove(evicted.SizeBytes) + q.logDrop(evicted.Kind, evicted.SizeBytes) + evictedCount++ + } + // Handle capacity-based eviction + q.nextGenID++ + if q.size == q.capacity { + evicted := q.data[q.head] + if q.activeGen == evicted.Generation { + q.activeGen = 0 + } + q.data[q.head] = QueueItem{} + q.head = (q.head + 1) % q.capacity + q.size-- + q.sizeBytes -= evicted.SizeBytes + QueueMemory.Remove(evicted.SizeBytes) + q.logDrop(evicted.Kind, evicted.SizeBytes) + } +} + +func (q *Queue) logDrop(kind string, sizeBytes int64) { + if q.Logger != nil { + q.Logger.PrintWarning("Dropped %s snapshot (%d bytes) from queue", kind, sizeBytes) + } + if q.DropCallback != nil { + q.DropCallback(kind, sizeBytes) + } +} diff --git a/state/queue_test.go b/state/queue_test.go new file mode 100644 index 000000000..6e7f93b3b --- /dev/null +++ b/state/queue_test.go @@ -0,0 +1,169 @@ +package state + +import ( + "bytes" + "context" + "fmt" + "sync" + "sync/atomic" + "testing" + "time" +) + +func TestQueue_FIFOEviction(t *testing.T) { + q := &Queue{data: make([]QueueItem, 5), capacity: 5} + q.cond = sync.NewCond(&q.mu) + ctx := context.Background() + // First item should be evicted after pusing 6 items into 5-capacity queue + for i := 0; i < 6; i++ { + q.PushBytes(fmt.Sprintf("%d", i), []byte{byte(i)}) + } + tx, err := q.Pop(ctx) + if err != nil { + t.Fatal(err) + } + if tx.Kind != "1" || !bytes.Equal(tx.Snapshot, []byte{1}) { + t.Errorf("expected item 1, got kind=%s snapshot=%v", tx.Kind, tx.Snapshot) + } + tx.Commit() + // Eviction during in-flight transaction is a no-op; second item remains + q2 := &Queue{data: make([]QueueItem, 2), capacity: 2} + q2.cond = sync.NewCond(&q2.mu) + q2.PushBytes("1", []byte("data")) + tx2, _ := q2.Pop(ctx) + q2.PushBytes("2", []byte("data")) // evicts tx1 + tx2.Commit() // safe no-op + tx3, err := q2.Pop(ctx) + if err != nil { + t.Fatal(err) + } + if tx3.Kind != "2" { + t.Errorf("expected item 2 after eviction, got %s", tx3.Kind) + } + tx3.Commit() +} + +func TestQueue_RollbackPreservesHead(t *testing.T) { + q := NewQueue(nil) + ctx := context.Background() + q.PushBytes("1", []byte("data")) + q.PushBytes("2", []byte("data")) + tx, _ := q.Pop(ctx) + tx.Rollback() + tx2, err := q.Pop(ctx) + if err != nil { + t.Fatal(err) + } + if tx2.Kind != "1" { + t.Errorf("expected Kind 1 after rollback, got %v", tx2.Kind) + } + tx2.Commit() +} + +func TestQueue_ConcurrentPushPopEviction(t *testing.T) { + q := &Queue{data: make([]QueueItem, 3), capacity: 3} + q.cond = sync.NewCond(&q.mu) + var dropCount int64 + q.DropCallback = func(string, int64) { atomic.AddInt64(&dropCount, 1) } + var wg sync.WaitGroup + for p := 0; p < 8; p++ { + wg.Add(1) + go func(id int) { + defer wg.Done() + for i := 0; i < 50; i++ { + q.PushBytes(fmt.Sprintf("p%d-i%d", id, i), []byte(fmt.Sprintf("data-%d-%d", id, i))) + } + }(p) + } + wg.Wait() + q.Close() + if dropCount == 0 { + t.Error("expected drops under capacity pressure") + } +} + +func TestQueue_CloseUnblocksPop(t *testing.T) { + q := NewQueue(nil) + ctx := context.Background() + var wg sync.WaitGroup + var errCount int64 + for i := 0; i < 10; i++ { + wg.Add(1) + go func() { + defer wg.Done() + _, err := q.Pop(ctx) + if err == nil { + t.Error("expected error from closed queue") + } else if err.Error() == "queue closed" { + atomic.AddInt64(&errCount, 1) + } + }() + } + go q.Close() + wg.Wait() + if atomic.LoadInt64(&errCount) != 10 { + t.Errorf("expected 10 'queue closed' errors, got %d", errCount) + } +} + +func TestQueue_ContextCancelUnblocksPop(t *testing.T) { + q := NewQueue(nil) + ctx, cancel := context.WithCancel(context.Background()) + var wg sync.WaitGroup + for i := 0; i < 5; i++ { + wg.Add(1) + go func() { + defer wg.Done() + q.Pop(ctx) + }() + } + cancel() + wg.Wait() + q.Close() +} + +func TestQueue_PushUnblockedByWaitingPop(t *testing.T) { + q := NewQueue(nil) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + var popDone sync.WaitGroup + popDone.Add(1) + go func() { + defer popDone.Done() + q.Pop(ctx) + }() + // Wait until Pop has entered cond.Wait() and released the mutex + for i := 0; i < 100; i++ { + q.mu.Lock() + empty := q.size == 0 + q.mu.Unlock() + if empty { + break + } + } + done := make(chan struct{}) + go func() { + q.PushBytes("test", []byte("data")) + close(done) + }() + select { + case <-done: + case <-time.After(2 * time.Second): + t.Fatal("Push blocked while Pop was waiting") + } + cancel() + popDone.Wait() +} + +func TestQueue_GlobalMemoryLimitEviction(t *testing.T) { + original := QueueMemory + defer func() { QueueMemory = original }() + QueueMemory = NewMemoryLimit(200) + q := NewQueue(nil) + for i := 0; i < 100; i++ { + q.PushBytes(fmt.Sprintf("item-%d", i), make([]byte, 10)) + } + if q.size > 25 { + t.Errorf("expected at most ~20 items, got %d", q.size) + } +} diff --git a/state/state.go b/state/state.go index 8e8b775ef..2b898dd2f 100644 --- a/state/state.go +++ b/state/state.go @@ -304,8 +304,7 @@ type Server struct { SelfTest *SelfTestResult - FullSnapshotUpload chan *pganalyze_collector.FullSnapshot - CompactSnapshotUpload chan *pganalyze_collector.CompactSnapshot + SnapshotQueue *Queue WebSocket *util.ReconnectingSocket InitialConfigReceived chan struct{} Pause atomic.Bool @@ -339,8 +338,7 @@ func MakeServer(config config.ServerConfig, testRun bool) *Server { ActivityStateMutex: &sync.Mutex{}, HighFreqStateMutex: &sync.Mutex{}, CollectionStatusMutex: &sync.Mutex{}, - FullSnapshotUpload: make(chan *pganalyze_collector.FullSnapshot), - CompactSnapshotUpload: make(chan *pganalyze_collector.CompactSnapshot), + SnapshotQueue: NewQueue(nil), InitialConfigReceived: make(chan struct{}, 1), QueryRuns: make(map[int64]*QueryRun), QueryRunsMutex: &sync.Mutex{}, diff --git a/util/reconnecting_socket.go b/util/reconnecting_socket.go index 59ebc3058..076f98229 100644 --- a/util/reconnecting_socket.go +++ b/util/reconnecting_socket.go @@ -10,10 +10,18 @@ import ( "github.com/gorilla/websocket" ) +// WriteRequest wraps data for a websocket write, including a response channel +// so the caller can learn whether the write succeeded. The Result channel +// receives the write error (nil on success). +type WriteRequest struct { + Data []byte + Result chan error +} + type ReconnectingSocket struct { // Channels shared with the caller Read chan []byte - Write chan []byte + Write chan WriteRequest // Initial arguments dialer websocket.Dialer @@ -38,7 +46,7 @@ var ErrorConnectRateLimited = errors.New("Skipping connection attempt because of func NewReconnectingSocket(ctx context.Context, logger *Logger, dialer websocket.Dialer, url string, headers map[string][]string, reconnectInterval time.Duration, clientErrorTimeout time.Duration) *ReconnectingSocket { w := &ReconnectingSocket{ Read: make(chan []byte), - Write: make(chan []byte), + Write: make(chan WriteRequest), ctx: ctx, dialer: dialer, url: url, @@ -156,8 +164,11 @@ func (w *ReconnectingSocket) connect(ctx context.Context) (int, error) { case <-connCtx.Done(): w.closeConnection() return - case data := <-w.Write: - err = conn.WriteMessage(websocket.BinaryMessage, data) + case req := <-w.Write: + err = conn.WriteMessage(websocket.BinaryMessage, req.Data) + if req.Result != nil { + req.Result <- err + } if err != nil { w.logger.PrintError("Error writing to websocket: %s", err) w.closeConnection() From ad899e3841e9bced05b9d7c43837e26cc1f17407 Mon Sep 17 00:00:00 2001 From: Sean Linsley Date: Thu, 23 Jul 2026 15:53:47 -0400 Subject: [PATCH 2/3] Always retain in-flight snapshots. This simplifies the code so generation tracking isn't needed --- state/queue.go | 92 ++++++++++++++++++--------------------------- state/queue_test.go | 8 ++-- 2 files changed, 40 insertions(+), 60 deletions(-) diff --git a/state/queue.go b/state/queue.go index 07d1e7248..58ed38bb4 100644 --- a/state/queue.go +++ b/state/queue.go @@ -13,14 +13,9 @@ import ( const DefaultCapacity = 500 -// Queue of snapshots ready for submission to pganalyze -// - Thread safe: Safe for concurrent use by multiple readers and writers. -// - Limited capacity: If full, Push drops the oldest item to make room. -// - Transactional: Pop locks the head item via a generation ID which is removed on Commit, and re-released on Rollback. -// - Generation tracking: Increments a counter on every Push to uniquely identify each item. -// If a Push evicts an in-flight item, its generation changes, safely ignoring late Commits or Rollbacks. -// - Memory-bounded: When the global QueueMemory limit is exceeded, Push drops the oldest items -// to stay within the budget. All queues share this single global limit. +// Queue of snapshots ready for submission to pganalyze. If snapshot +// submission fails, the original snapshot order is retained up until +// snapshots must be dropped to stay within the capacity and memory limits. type Queue struct { mu sync.Mutex cond *sync.Cond @@ -30,9 +25,8 @@ type Queue struct { tail int size int sizeBytes int64 - activeGen uint64 + inFlight bool // true while an item is being uploaded (between Pop and Commit/Rollback) closed bool - nextGenID uint64 Logger *util.Logger DropCallback func(kind string, sizeBytes int64) } @@ -47,10 +41,9 @@ func NewQueue(logger *util.Logger) *Queue { return q } -// Push adds an item to the tail. If full, it drops the oldest item to make room. -// When the global queue memory limit is exceeded, Push drops the oldest items -// (starting with the current queue's oldest) to stay within the budget. -// All queues share the single global memory limit defined by QueueMemory. +// Push adds an item to the tail. If full or over the global memory limit, +// it drops the oldest items to make room, unless that item is currently +// being uploaded (in-flight). In that case, the new push is dropped instead. func (q *Queue) Push(kind string, snapshot proto.Message) (err error) { data, err := proto.Marshal(snapshot) if err != nil { @@ -63,7 +56,7 @@ func (q *Queue) Push(kind string, snapshot proto.Message) (err error) { return q.PushBytes(kind, buf.Bytes()) } -// Factored out from Push so it can be called by tests without constructing real snapshots +// Factored out from Push so it can be called by tests without building real snapshots func (q *Queue) PushBytes(kind string, bytes []byte) (err error) { sizeBytes := int64(len(bytes)) q.mu.Lock() @@ -73,10 +66,9 @@ func (q *Queue) PushBytes(kind string, bytes []byte) (err error) { } q.makeSpace(sizeBytes) q.data[q.tail] = QueueItem{ - Kind: kind, - Snapshot: bytes, - SizeBytes: sizeBytes, - Generation: q.nextGenID, + Kind: kind, + Snapshot: bytes, + SizeBytes: sizeBytes, } q.tail = (q.tail + 1) % q.capacity q.size++ @@ -86,13 +78,11 @@ func (q *Queue) PushBytes(kind string, bytes []byte) (err error) { } // Pop blocks until an item is ready or the queue closes. -// On success, it locks the head item and returns a Transaction handle. -// -// The caller is woken by Push (Signal), Commit/Rollback (Broadcast), or Close (Broadcast). +// On success, it marks the head as in-flight and returns a Transaction handle. func (q *Queue) Pop(ctx context.Context) (*Transaction, error) { q.mu.Lock() defer q.mu.Unlock() - for (q.size == 0 || q.activeGen != 0) && !q.closed && ctx.Err() == nil { + for (q.size == 0 || q.inFlight) && !q.closed && ctx.Err() == nil { q.cond.Wait() } if ctx.Err() != nil { @@ -102,13 +92,12 @@ func (q *Queue) Pop(ctx context.Context) (*Transaction, error) { return nil, errors.New("queue closed") } item := q.data[q.head] - q.activeGen = item.Generation + q.inFlight = true return &Transaction{ - Kind: item.Kind, - Snapshot: item.Snapshot, - SizeBytes: item.SizeBytes, - generation: item.Generation, - q: q, + Kind: item.Kind, + Snapshot: item.Snapshot, + SizeBytes: item.SizeBytes, + q: q, }, nil } @@ -123,25 +112,23 @@ func (q *Queue) Close() { } type QueueItem struct { - Kind string - Snapshot []byte - SizeBytes int64 - Generation uint64 + Kind string + Snapshot []byte + SizeBytes int64 } type Transaction struct { - Kind string - Snapshot []byte - SizeBytes int64 - generation uint64 - q *Queue + Kind string + Snapshot []byte + SizeBytes int64 + q *Queue } +// Commit advances past the head item, marking it as successfully uploaded. func (t *Transaction) Commit() { t.q.mu.Lock() defer t.q.mu.Unlock() - // Validate that the transaction hasn't been evicted or superseded - if t.q.activeGen != t.generation || t.q.data[t.q.head].Generation != t.generation { + if !t.q.inFlight { return } t.q.data[t.q.head] = QueueItem{} @@ -149,37 +136,34 @@ func (t *Transaction) Commit() { t.q.size-- t.q.sizeBytes -= t.SizeBytes QueueMemory.Remove(t.SizeBytes) - t.q.activeGen = 0 + t.q.inFlight = false t.q.cond.Broadcast() } +// Rollback leaves the head item in place so it will be retried on the next Pop. func (t *Transaction) Rollback() { t.q.mu.Lock() defer t.q.mu.Unlock() - // Validate that the transaction hasn't been evicted or superseded - if t.q.activeGen != t.generation || t.q.data[t.q.head].Generation != t.generation { + if !t.q.inFlight { return } - t.q.activeGen = 0 + t.q.inFlight = false t.q.cond.Broadcast() } +// Does nothing if the head snapshot is in-flight to preserve ordering on retry. func (q *Queue) makeSpace(sizeBytes int64) { QueueMemory.Add(sizeBytes) - // Evict items to satisfy the global memory limit + // Evict items to satisfy the global memory limit. evictedCount := 0 for evictedCount <= 100 { if !QueueMemory.OverLimit() { break } - if q.size == 0 { - // This queue is empty; another server's queue is the problem + if q.size == 0 || q.inFlight { break } evicted := q.data[q.head] - if q.activeGen == evicted.Generation { - q.activeGen = 0 - } q.data[q.head] = QueueItem{} q.head = (q.head + 1) % q.capacity q.size-- @@ -188,13 +172,9 @@ func (q *Queue) makeSpace(sizeBytes int64) { q.logDrop(evicted.Kind, evicted.SizeBytes) evictedCount++ } - // Handle capacity-based eviction - q.nextGenID++ - if q.size == q.capacity { + // Handle capacity-based eviction. + if q.size == q.capacity && !q.inFlight { evicted := q.data[q.head] - if q.activeGen == evicted.Generation { - q.activeGen = 0 - } q.data[q.head] = QueueItem{} q.head = (q.head + 1) % q.capacity q.size-- diff --git a/state/queue_test.go b/state/queue_test.go index 6e7f93b3b..615595fee 100644 --- a/state/queue_test.go +++ b/state/queue_test.go @@ -26,19 +26,19 @@ func TestQueue_FIFOEviction(t *testing.T) { t.Errorf("expected item 1, got kind=%s snapshot=%v", tx.Kind, tx.Snapshot) } tx.Commit() - // Eviction during in-flight transaction is a no-op; second item remains + // In-flight items are never evicted; they stay at head until committed. q2 := &Queue{data: make([]QueueItem, 2), capacity: 2} q2.cond = sync.NewCond(&q2.mu) q2.PushBytes("1", []byte("data")) tx2, _ := q2.Pop(ctx) - q2.PushBytes("2", []byte("data")) // evicts tx1 - tx2.Commit() // safe no-op + q2.PushBytes("2", []byte("data")) // fits without eviction (size 1 < capacity 2) + tx2.Commit() // removes "1" tx3, err := q2.Pop(ctx) if err != nil { t.Fatal(err) } if tx3.Kind != "2" { - t.Errorf("expected item 2 after eviction, got %s", tx3.Kind) + t.Errorf("expected item 2 after commit, got %s", tx3.Kind) } tx3.Commit() } From 7370d935577cef925917ceae3d7fcf4371e71c13 Mon Sep 17 00:00:00 2001 From: Sean Linsley Date: Thu, 23 Jul 2026 16:13:34 -0400 Subject: [PATCH 3/3] Log when PushBytes must drop new items, fix backoff delay in upload loop --- output/upload.go | 2 +- state/queue.go | 7 +++++++ 2 files changed, 8 insertions(+), 1 deletion(-) diff --git a/output/upload.go b/output/upload.go index 5875bc2ee..6a80f62fc 100644 --- a/output/upload.go +++ b/output/upload.go @@ -30,7 +30,7 @@ func snapshotUploadForServer(ctx context.Context, server *state.Server, logger * for { if failed { - delay = min(5*delay, 10*time.Second) // Increasing backoff delay in case of failure + delay = min(delay*5+10*time.Millisecond, 10*time.Second) } else { delay = 10 * time.Millisecond // Small delay to avoid high CPU usage in loop } diff --git a/state/queue.go b/state/queue.go index 58ed38bb4..66e7b8046 100644 --- a/state/queue.go +++ b/state/queue.go @@ -64,7 +64,14 @@ func (q *Queue) PushBytes(kind string, bytes []byte) (err error) { if q.closed { return } + beforeSize := q.size q.makeSpace(sizeBytes) + // If makeSpace couldn't evict (e.g., head is in-flight), drop the new item. + if beforeSize == q.size && beforeSize >= q.capacity { + QueueMemory.Remove(sizeBytes) + q.logDrop(kind, sizeBytes) + return + } q.data[q.tail] = QueueItem{ Kind: kind, Snapshot: bytes,