warpbox-dev/backend/libs/handlers/upload_group.go

package handlers

import (
	"sync"
	"time"
)

// uploadGroupWindow is how long after a batched upload a follow-up upload with
// the same X-Warpbox-Batch value (and same account/IP) is folded into the same
// box. ShareX sends a multi-file selection as separate back-to-back requests;
// the batch header lets it land them in one box.
const uploadGroupWindow = 20 * time.Second

// uploadBatchHeader is the opt-in request header. Without it, uploads behave
// exactly as before (one box per request). With it, requests sharing the same
// value (per account/IP) within uploadGroupWindow are grouped into one box.
const uploadBatchHeader = "X-Warpbox-Batch"

// uploadGroupPruneInterval is how often entryFor drops stale entries so the map
// can't grow without bound (one key per account/IP + batch value otherwise).
const uploadGroupPruneInterval = 5 * time.Minute

// uploadGrouper tracks the most recent box per batch key so opt-in batched
// uploads land in a single box. Each key has its own lock, which also serialises
// that key's concurrent uploads so they append to the same box instead of racing
// to create several.
type uploadGrouper struct {
	mu        sync.Mutex
	entries   map[string]*uploadGroupEntry
	lastPrune time.Time
}

type uploadGroupEntry struct {
	mu        sync.Mutex
	boxID     string
	manageURL string
	deleteURL string
	at        time.Time
}

func newUploadGrouper() *uploadGrouper {
	return &uploadGrouper{entries: make(map[string]*uploadGroupEntry)}
}

func (g *uploadGrouper) entryFor(key string) *uploadGroupEntry {
	g.mu.Lock()
	defer g.mu.Unlock()
	g.pruneLocked(time.Now())
	entry, ok := g.entries[key]
	if !ok {
		entry = &uploadGroupEntry{}
		g.entries[key] = entry
	}
	return entry
}

// pruneLocked drops entries whose last use is well past the grouping window so
// the map stays bounded to recently-active keys. Callers must hold g.mu. Entries
// currently in use, or freshly created but not yet used (zero timestamp), are
// kept to avoid removing one a request is about to populate.
func (g *uploadGrouper) pruneLocked(now time.Time) {
	if now.Sub(g.lastPrune) < uploadGroupPruneInterval {
		return
	}
	g.lastPrune = now
	for key, entry := range g.entries {
		if !entry.mu.TryLock() {
			continue
		}
		stale := !entry.at.IsZero() && now.Sub(entry.at) > 2*uploadGroupWindow
		entry.mu.Unlock()
		if stale {
			delete(g.entries, key)
		}
	}
}