archivmail/internal/storage/quota.go

package storage

import (
	"context"
	"errors"
	"sync"
	"time"
)

// ErrQuotaExceeded is returned when a tenant's storage or email quota is exceeded.
var ErrQuotaExceeded = errors.New("storage: tenant quota exceeded")

const quotaCacheTTL = 60 * time.Second

type quotaEntry struct {
	exceeded bool
	expiry   time.Time
}

type quotaCache struct {
	mu      sync.Mutex
	entries map[int64]quotaEntry
}

var qCache = &quotaCache{entries: make(map[int64]quotaEntry)}

// CheckQuota checks whether a tenant has exceeded its storage or email quota.
// Returns ErrQuotaExceeded if exceeded. Results are cached for 60 seconds.
// If tenantID is nil (no-tenant mail), quota is not enforced.
func (s *Store) CheckQuota(ctx context.Context, tenantID *int64) error {
	if s.db == nil || tenantID == nil {
		return nil
	}

	id := *tenantID

	// Check cache first
	qCache.mu.Lock()
	if e, ok := qCache.entries[id]; ok && time.Now().Before(e.expiry) {
		exceeded := e.exceeded
		qCache.mu.Unlock()
		if exceeded {
			return ErrQuotaExceeded
		}
		return nil
	}
	qCache.mu.Unlock()

	// Query quota and current usage in one statement
	var maxStorageBytes *int64
	var maxEmails *int64
	var currentBytes int64
	var currentEmails int64

	err := s.db.QueryRow(ctx, `
		SELECT t.max_storage_bytes, t.max_emails,
		       COALESCE(SUM(e.size), 0) AS total_bytes,
		       COUNT(e.id)              AS total_emails
		FROM tenants t
		LEFT JOIN emails e ON e.tenant_id = t.id
		WHERE t.id = $1
		GROUP BY t.id
	`, id).Scan(&maxStorageBytes, &maxEmails, &currentBytes, &currentEmails)
	if err != nil {
		// Tenant not found or DB error — allow the mail to pass through
		return nil
	}

	exceeded := false
	if maxStorageBytes != nil && currentBytes >= *maxStorageBytes {
		exceeded = true
	}
	if maxEmails != nil && currentEmails >= *maxEmails {
		exceeded = true
	}

	// Write to cache
	qCache.mu.Lock()
	qCache.entries[id] = quotaEntry{exceeded: exceeded, expiry: time.Now().Add(quotaCacheTTL)}
	qCache.mu.Unlock()

	if exceeded {
		return ErrQuotaExceeded
	}
	return nil
}

// InvalidateQuotaCache removes the cached quota entry for a tenant.
// Call this after quota changes to ensure fresh values on the next check.
func InvalidateQuotaCache(tenantID int64) {
	qCache.mu.Lock()
	delete(qCache.entries, tenantID)
	qCache.mu.Unlock()
}