zygostates/internal/zygostates/uploader.go

package zygostates

import (
	"bytes"
	"compress/gzip"
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"net/url"
	"path"
	"strconv"
	"sync"
	"time"
)

type PutRequest struct {
	Bucket          string
	Key             string
	Body            []byte
	ContentType     string
	ContentEncoding string
	Metadata        map[string]string
}

type ObjectStore interface {
	Put(context.Context, PutRequest) error
}

type UploadedRange struct {
	Path       string `json:"path"`
	Generation int    `json:"generation"`
	Start      int64  `json:"start"`
	End        int64  `json:"end"`
	SHA256     string `json:"sha256"`
	Key        string `json:"key"`
}

type UploadTick struct {
	Timestamp time.Time       `json:"timestamp"`
	Sequence  uint64          `json:"sequence"`
	Uploaded  []UploadedRange `json:"uploaded"`
}

type uploadPosition struct {
	offset     int64
	generation int
}

type Uploader struct {
	store     ObjectStore
	artifacts *ArtifactStore
	bucket    string
	baseKey   string
	positions map[string]uploadPosition
	sequence  uint64
	maxChunk  int
	mu        sync.Mutex
}

func newUploader(store ObjectStore, artifacts *ArtifactStore, config UploadConfig, identity, runID string) *Uploader {
	return &Uploader{
		store: store, artifacts: artifacts, bucket: config.Bucket,
		baseKey:   path.Join(config.Prefix, identity, runID),
		positions: make(map[string]uploadPosition), maxChunk: 8 * 1024 * 1024,
	}
}

func (u *Uploader) Upload(ctx context.Context) (int, error) {
	u.mu.Lock()
	defer u.mu.Unlock()
	snapshots, err := u.artifacts.Snapshots()
	if err != nil {
		return 0, err
	}
	staged := make(map[string]uploadPosition, len(u.positions))
	for name, position := range u.positions {
		staged[name] = position
	}
	var uploaded []UploadedRange
	for _, snapshot := range snapshots {
		position := staged[snapshot.Name]
		if int64(len(snapshot.Contents)) < position.offset {
			position.offset = 0
			position.generation++
		}
		for position.offset < int64(len(snapshot.Contents)) {
			end := min(position.offset+int64(u.maxChunk), int64(len(snapshot.Contents)))
			content := snapshot.Contents[position.offset:end]
			digest := sha256.Sum256(content)
			digestText := hex.EncodeToString(digest[:])
			key := path.Join(u.baseKey, "files", url.PathEscape(snapshot.Name), fmt.Sprintf("%04d", position.generation), fmt.Sprintf("%020d-%020d.gz", position.offset, end))
			var compressed bytes.Buffer
			writer, _ := gzip.NewWriterLevel(&compressed, gzip.BestSpeed)
			if _, err := writer.Write(content); err != nil {
				return len(uploaded), err
			}
			if err := writer.Close(); err != nil {
				return len(uploaded), err
			}
			if err := u.store.Put(ctx, PutRequest{
				Bucket: u.bucket, Key: key, Body: compressed.Bytes(), ContentType: "application/octet-stream", ContentEncoding: "gzip",
				Metadata: map[string]string{"source-path": snapshot.Name, "generation": strconv.Itoa(position.generation), "start": strconv.FormatInt(position.offset, 10), "end": strconv.FormatInt(end, 10), "sha256": digestText},
			}); err != nil {
				return len(uploaded), err
			}
			uploaded = append(uploaded, UploadedRange{Path: snapshot.Name, Generation: position.generation, Start: position.offset, End: end, SHA256: digestText, Key: key})
			position.offset = end
		}
		staged[snapshot.Name] = position
	}
	tick := UploadTick{Timestamp: time.Now().UTC(), Sequence: u.sequence, Uploaded: uploaded}
	tickBody, err := json.Marshal(tick)
	if err != nil {
		return len(uploaded), err
	}
	tickKey := path.Join(u.baseKey, "ticks", fmt.Sprintf("%08d.json", u.sequence))
	if err := u.store.Put(ctx, PutRequest{Bucket: u.bucket, Key: tickKey, Body: tickBody, ContentType: "application/json"}); err != nil {
		return len(uploaded), err
	}
	u.positions = staged
	u.sequence++
	return len(uploaded), nil
}

type PeriodicUploader struct {
	uploader *Uploader
	interval time.Duration
	store    *ArtifactStore
}

func (u *PeriodicUploader) Start(ctx context.Context) *sync.WaitGroup {
	var group sync.WaitGroup
	group.Add(1)
	go func() {
		defer group.Done()
		ticker := time.NewTicker(u.interval)
		defer ticker.Stop()
		for {
			select {
			case <-ctx.Done():
				return
			case <-ticker.C:
				uploadContext, cancel := context.WithTimeout(context.WithoutCancel(ctx), 30*time.Second)
				count, err := u.uploader.Upload(uploadContext)
				cancel()
				fields := map[string]any{"uploaded_ranges": count}
				if err != nil {
					fields["error"] = err.Error()
				}
				_ = u.store.AppendEvent(Event{Timestamp: time.Now().UTC(), Type: "upload", Fields: fields})
			}
		}
	}()
	return &group
}