Retry failed snapshot submission

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 {

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 {

output/upload.go

@@ -1,8 +1,6 @@
 package output
 
 import (
-	"bytes"
-	"compress/zlib"
 	"context"
 	"errors"
 	"fmt"
@@ -11,55 +9,58 @@ 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) {
 	if opts.ForceEmptyGrant {
 		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(delay*5+10*time.Millisecond, 10*time.Second)
+		} 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)
 }

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,110 +32,37 @@ 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)
 	if err != nil {
 		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

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

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
+}