go-cart-actor/pkg/actor/disk_storage.go

package actor

import (
	"context"
	"errors"
	"fmt"
	"log"
	"os"
	"path/filepath"
	"sync"
	"time"

	"google.golang.org/protobuf/proto"
)

const defaultWriterIdleTTL = 5 * time.Minute

type QueueEvent struct {
	TimeStamp time.Time
	Message   proto.Message
}

// diskLogWriter holds one append-only log file kept open between writes.
type diskLogWriter struct {
	mu       sync.Mutex
	file     *os.File
	lastUsed time.Time
	purged   bool // set by purgeIdleWriters before closing; getWriter re-checks after acquiring mu
}

type DiskStorage[V any] struct {
	*StateStorage
	path          string
	done          chan struct{}
	queue         *sync.Map // map[uint64][]QueueEvent
	dirOnce       sync.Once
	dirErr        error
	writersMu     sync.Mutex
	writers       map[uint64]*diskLogWriter
	writerIdleTTL time.Duration
}

type LogStorage[V any] interface {
	LoadEvents(ctx context.Context, id uint64, grain Grain[V]) error
	LoadEventsFunc(ctx context.Context, id uint64, grain Grain[V], condition func(msg proto.Message, index int, timeStamp time.Time) bool) error
	AppendMutations(id uint64, msg ...proto.Message) error
}

func NewDiskStorage[V any](path string, registry MutationRegistry) *DiskStorage[V] {
	s := &DiskStorage[V]{
		StateStorage:  NewState(registry),
		path:          path,
		done:          make(chan struct{}),
		writers:       make(map[uint64]*diskLogWriter),
		writerIdleTTL: defaultWriterIdleTTL,
	}
	go s.purgeLoop()
	return s
}

func (s *DiskStorage[V]) ensureDir() error {
	s.dirOnce.Do(func() {
		s.dirErr = os.MkdirAll(s.path, 0o755)
	})
	return s.dirErr
}

// getWriter returns an open append handle for id. The caller must unlock w.mu
// when finished writing.


func (s *DiskStorage[V]) writeEvents(id uint64, events []QueueEvent) error {
	if len(events) == 0 {
		return nil
	}
	w, err := s.getWriter(id)
	if err != nil {
		return err
	}
	defer w.mu.Unlock()

	for _, evt := range events {
		if err := s.Append(w.file, evt.Message, evt.TimeStamp); err != nil {
			return err
		}
	}
	w.lastUsed = time.Now()
	return nil
}

func (s *DiskStorage[V]) writeMutations(id uint64, msgs ...proto.Message) error {
	if len(msgs) == 0 {
		return nil
	}
	w, err := s.getWriter(id)
	if err != nil {
		return err
	}
	defer w.mu.Unlock()

	now := time.Now()
	for _, m := range msgs {
		if err := s.Append(w.file, m, now); err != nil {
			return err
		}
	}
	w.lastUsed = time.Now()
	return nil
}

func (s *DiskStorage[V]) purgeLoop() {
	ticker := time.NewTicker(time.Minute)
	defer ticker.Stop()
	for {
		select {
		case <-s.done:
			return
		case <-ticker.C:
			s.purgeIdleWriters()
		}
	}
}

func (s *DiskStorage[V]) purgeIdleWriters() {
	cutoff := time.Now().Add(-s.writerIdleTTL)

	s.writersMu.Lock()
	defer s.writersMu.Unlock()

	for id, w := range s.writers {
		w.mu.Lock()
		if w.lastUsed.Before(cutoff) {
			w.purged = true
			delete(s.writers, id)
			if err := w.file.Close(); err != nil {
				log.Printf("failed to close idle event log %d: %v", id, err)
			}
		}
		w.mu.Unlock()
	}
}

func (s *DiskStorage[V]) getWriter(id uint64) (*diskLogWriter, error) {
	if err := s.ensureDir(); err != nil {
		return nil, err
	}

	for {
		s.writersMu.Lock()
		w, ok := s.writers[id]
		if !ok {
			fh, err := os.OpenFile(s.logPath(id), os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0o644)
			if err != nil {
				s.writersMu.Unlock()
				return nil, err
			}
			w = &diskLogWriter{file: fh, lastUsed: time.Now()}
			s.writers[id] = w
		}
		s.writersMu.Unlock()

		w.mu.Lock()
		if !w.purged {
			w.lastUsed = time.Now()
			return w, nil
		}
		// Writer was purged by purgeIdleWriters between our map lookup and acquiring w.mu.
		// Release it and retry — the map no longer contains this entry, so the next
		// iteration will create a fresh writer.
		w.mu.Unlock()
	}
}

func (s *DiskStorage[V]) closeWriters() {
	s.writersMu.Lock()
	defer s.writersMu.Unlock()
	for id, w := range s.writers {
		w.mu.Lock()
		w.purged = true
		if err := w.file.Close(); err != nil {
			log.Printf("failed to close event log %d: %v", id, err)
		}
		delete(s.writers, id)
		w.mu.Unlock()
	}
}

func (s *DiskStorage[V]) SaveLoop(duration time.Duration) {
	s.queue = &sync.Map{}
	ticker := time.NewTicker(duration)
	defer ticker.Stop()
	for {
		select {
		case <-s.done:
			s.save()
			return
		case <-ticker.C:
			s.save()
		}
	}
}

func (s *DiskStorage[V]) save() {
	carts := 0
	lines := 0
	s.queue.Range(func(key, value any) bool {
		id := key.(uint64)
		qe, ok := value.([]QueueEvent)
		if !ok {
			s.queue.Delete(id)
			return true
		}
		if err := s.writeEvents(id, qe); err != nil {
			log.Printf("failed to append events for grain %d: %v", id, err)
			return true
		}
		lines += len(qe)
		carts++
		s.queue.Delete(id)
		return true
	})
	if lines > 0 {
		log.Printf("Appended %d carts and %d lines to disk", carts, lines)
	}
}

func (s *DiskStorage[V]) logPath(id uint64) string {
	return filepath.Join(s.path, fmt.Sprintf("%d.events.log", id))
}

func (s *DiskStorage[V]) LoadEventsFunc(ctx context.Context, id uint64, grain Grain[V], condition func(msg proto.Message, index int, timeStamp time.Time) bool) error {
	path := s.logPath(id)
	if _, err := os.Stat(path); errors.Is(err, os.ErrNotExist) {
		// No log -> nothing to replay
		return nil
	}

	fh, err := os.Open(path)
	if err != nil {
		return fmt.Errorf("open replay file: %w", err)
	}
	defer fh.Close()
	index := 0
	return s.Load(fh, func(msg proto.Message, when time.Time) {
		if condition(msg, index, when) {
			s.registry.Apply(ctx, grain, msg)
		}
		index++
	})
}

func (s *DiskStorage[V]) LoadEvents(ctx context.Context, id uint64, grain Grain[V]) error {
	path := s.logPath(id)
	if _, err := os.Stat(path); errors.Is(err, os.ErrNotExist) {
		// No log -> nothing to replay
		return nil
	}

	fh, err := os.Open(path)
	if err != nil {
		return fmt.Errorf("open replay file: %w", err)
	}
	defer fh.Close()
	return s.Load(fh, func(msg proto.Message, _ time.Time) {
		s.registry.Apply(ctx, grain, msg)
	})
}

func (s *DiskStorage[V]) Close() {
	if s.queue != nil {
		s.save()
	}
	s.closeWriters()
	close(s.done)
}

func (s *DiskStorage[V]) AppendMutations(id uint64, msg ...proto.Message) error {
	if s.queue != nil {
		queue := make([]QueueEvent, 0)
		data, found := s.queue.Load(id)
		if found {
			queue = data.([]QueueEvent)
		}
		for _, m := range msg {
			queue = append(queue, QueueEvent{Message: m, TimeStamp: time.Now()})
		}
		s.queue.Store(id, queue)
		return nil
	}
	return s.writeMutations(id, msg...)
}

func (s *DiskStorage[V]) openWriterCount() int {
	s.writersMu.Lock()
	defer s.writersMu.Unlock()
	return len(s.writers)
}