idempotency.go

package broker

import (
	"fmt"
	"sync"
	"time"
)

type ProducerSession struct {
	ProducerID string
	Epoch      int32
	CreatedAt  time.Time
	LastActive time.Time
	Sequences  map[int]int64 // partition -> last sequence number
}

type IdempotencyConfig struct {
	SessionTimeout time.Duration
	MaxProducers   int
	Enabled        bool
}

func DefaultIdempotencyConfig() IdempotencyConfig {
	return IdempotencyConfig{
		SessionTimeout: 15 * time.Minute,
		MaxProducers:   10000,
		Enabled:        true,
	}
}

type IdempotencyManager struct {
	mu       sync.RWMutex
	config   IdempotencyConfig
	sessions map[string]*ProducerSession
	nextID   int64
	stopCh   chan struct{}

	duplicatesRejected int64
	sessionsCreated    int64
	sessionsExpired    int64
}

func NewIdempotencyManager(config IdempotencyConfig) *IdempotencyManager {
	im := &IdempotencyManager{
		config:   config,
		sessions: make(map[string]*ProducerSession),
		nextID:   time.Now().UnixNano(),
		stopCh:   make(chan struct{}),
	}

	go im.cleanupLoop()

	return im
}

// InitProducer registers a producer and returns its ID and epoch.
// If the producer already exists, the epoch is bumped to fence the old instance.
func (im *IdempotencyManager) InitProducer(producerID string) (string, int32, error) {
	im.mu.Lock()
	defer im.mu.Unlock()

	if !im.config.Enabled {
		return "", 0, nil
	}

	if session, exists := im.sessions[producerID]; exists {
		session.Epoch++
		session.LastActive = time.Now()
		session.Sequences = make(map[int]int64)
		return session.ProducerID, session.Epoch, nil
	}

	if producerID == "" {
		im.nextID++
		producerID = fmt.Sprintf("producer-%d", im.nextID)
	}

	if len(im.sessions) >= im.config.MaxProducers {
		im.evictOldest()
	}

	session := &ProducerSession{
		ProducerID: producerID,
		Epoch:      0,
		CreatedAt:  time.Now(),
		LastActive: time.Now(),
		Sequences:  make(map[int]int64),
	}
	im.sessions[producerID] = session
	im.sessionsCreated++

	return producerID, 0, nil
}

// CheckAndUpdate validates a produce request for idempotency.
// Returns ErrDuplicateSequence for duplicates, ErrInvalidProducerEpoch for stale epochs.
func (im *IdempotencyManager) CheckAndUpdate(producerID string, epoch int32, partition int, sequence int64) error {
	im.mu.Lock()
	defer im.mu.Unlock()

	if !im.config.Enabled {
		return nil
	}

	session, exists := im.sessions[producerID]
	if !exists {
		// expired or unknown session — allow but don't track
		return nil
	}

	if epoch < session.Epoch {
		return ErrInvalidProducerEpoch
	}

	lastSeq, hasSeq := session.Sequences[partition]
	if hasSeq && sequence <= lastSeq {
		im.duplicatesRejected++
		return ErrDuplicateSequence
	}

	session.Sequences[partition] = sequence
	session.LastActive = time.Now()

	return nil
}

func (im *IdempotencyManager) GetSession(producerID string) (*ProducerSession, bool) {
	im.mu.RLock()
	defer im.mu.RUnlock()

	session, exists := im.sessions[producerID]
	if !exists {
		return nil, false
	}
	cp := *session
	cp.Sequences = make(map[int]int64)
	for k, v := range session.Sequences {
		cp.Sequences[k] = v
	}
	return &cp, true
}

func (im *IdempotencyManager) ExpireSession(producerID string) {
	im.mu.Lock()
	defer im.mu.Unlock()

	delete(im.sessions, producerID)
	im.sessionsExpired++
}

func (im *IdempotencyManager) GetStats() (activeSessions, duplicatesRejected, sessionsCreated, sessionsExpired int64) {
	im.mu.RLock()
	defer im.mu.RUnlock()

	return int64(len(im.sessions)), im.duplicatesRejected, im.sessionsCreated, im.sessionsExpired
}

// evictOldest removes the least-recently-active session. Must hold im.mu.
func (im *IdempotencyManager) evictOldest() {
	var oldest *ProducerSession
	var oldestID string

	for id, session := range im.sessions {
		if oldest == nil || session.LastActive.Before(oldest.LastActive) {
			oldest = session
			oldestID = id
		}
	}

	if oldestID != "" {
		delete(im.sessions, oldestID)
		im.sessionsExpired++
	}
}

func (im *IdempotencyManager) cleanupLoop() {
	ticker := time.NewTicker(1 * time.Minute)
	defer ticker.Stop()

	for {
		select {
		case <-ticker.C:
			im.cleanupExpired()
		case <-im.stopCh:
			return
		}
	}
}

func (im *IdempotencyManager) Stop() {
	close(im.stopCh)
}

func (im *IdempotencyManager) cleanupExpired() {
	im.mu.Lock()
	defer im.mu.Unlock()

	cutoff := time.Now().Add(-im.config.SessionTimeout)
	for id, session := range im.sessions {
		if session.LastActive.Before(cutoff) {
			delete(im.sessions, id)
			im.sessionsExpired++
		}
	}
}

type IdempotencySnapshot struct {
	Sessions []ProducerSessionSnapshot `json:"sessions"`
}

type ProducerSessionSnapshot struct {
	ProducerID string        `json:"producer_id"`
	Epoch      int32         `json:"epoch"`
	Sequences  map[int]int64 `json:"sequences"`
	LastActive int64         `json:"last_active"`
}

func (im *IdempotencyManager) Snapshot() IdempotencySnapshot {
	im.mu.RLock()
	defer im.mu.RUnlock()

	snap := IdempotencySnapshot{
		Sessions: make([]ProducerSessionSnapshot, 0, len(im.sessions)),
	}

	for _, session := range im.sessions {
		seqs := make(map[int]int64)
		for k, v := range session.Sequences {
			seqs[k] = v
		}
		snap.Sessions = append(snap.Sessions, ProducerSessionSnapshot{
			ProducerID: session.ProducerID,
			Epoch:      session.Epoch,
			Sequences:  seqs,
			LastActive: session.LastActive.UnixMilli(),
		})
	}

	return snap
}

func (im *IdempotencyManager) Restore(snap IdempotencySnapshot) {
	im.mu.Lock()
	defer im.mu.Unlock()

	im.sessions = make(map[string]*ProducerSession)
	for _, ss := range snap.Sessions {
		seqs := make(map[int]int64)
		for k, v := range ss.Sequences {
			seqs[k] = v
		}
		im.sessions[ss.ProducerID] = &ProducerSession{
			ProducerID: ss.ProducerID,
			Epoch:      ss.Epoch,
			Sequences:  seqs,
			LastActive: time.UnixMilli(ss.LastActive),
			CreatedAt:  time.UnixMilli(ss.LastActive),
		}
	}
}