ITWorx-Pulse-Public/internal/alertworker/memory_store.go

package alertworker

import (
	"context"
	"sync"
	"time"

	"github.com/itworx/pulse/internal/alert"
)

type memoryLease struct {
	lease      Lease
	status     string
	leaseUntil time.Time
}

type MemoryLeaseStore struct {
	mu      sync.Mutex
	records map[string]memoryLease
}

func NewMemoryLeaseStore() *MemoryLeaseStore {
	return &MemoryLeaseStore{records: make(map[string]memoryLease)}
}

func (s *MemoryLeaseStore) Acquire(_ context.Context, jobType, jobKey string, scheduledAt time.Time, owner string, now time.Time, ttl time.Duration) (Lease, bool, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.records == nil {
		s.records = make(map[string]memoryLease)
	}
	key := leaseKey(jobType, jobKey, scheduledAt)
	if record, ok := s.records[key]; ok {
		if record.status != "running" || record.leaseUntil.After(now) {
			return Lease{}, false, nil
		}
	}
	lease := Lease{ID: alert.NewID(), JobKey: jobKey, ScheduledAt: scheduledAt.UTC(), Owner: owner}
	s.records[key] = memoryLease{lease: lease, status: "running", leaseUntil: now.Add(ttl)}
	return lease, true, nil
}

func (s *MemoryLeaseStore) Complete(_ context.Context, lease Lease, status, _ string) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	key := leaseKey(jobType, lease.JobKey, lease.ScheduledAt)
	record, ok := s.records[key]
	if !ok || record.lease.ID != lease.ID || record.lease.Owner != lease.Owner {
		return ErrLeaseLost
	}
	record.status = status
	record.leaseUntil = time.Time{}
	s.records[key] = record
	return nil
}

func (s *MemoryLeaseStore) Status(jobKey string, scheduledAt time.Time) string {
	s.mu.Lock()
	defer s.mu.Unlock()
	record, ok := s.records[leaseKey(jobType, jobKey, scheduledAt)]
	if !ok {
		return ""
	}
	return record.status
}

func leaseKey(jobType, jobKey string, scheduledAt time.Time) string {
	return jobType + "|" + jobKey + "|" + scheduledAt.UTC().Format(time.RFC3339Nano)
}