fix(sweeper): move the scheduled task map to LiveStore port

internal/core/application/service.go

@@ -197,7 +197,7 @@ func NewService(
 		cache:                     cache,
 		scanner:                   scanner,
 		sweeper: newSweeper(
-			wallet, repoManager, builder, scheduler, noteUriPrefix,
+			wallet, repoManager, builder, scheduler, noteUriPrefix, cache,
 		),
 		boardingExitDelay:             boardingExitDelay,
 		operatorPrvkey:                operatorSigningKey,
@@ -3376,8 +3376,13 @@ func (s *service) listenToScannerNotifications() {
 							// remove sweeper task for the associated checkpoint outputs
 							for _, in := range ptx.UnsignedTx.TxIn {
 								taskId := in.PreviousOutPoint.Hash.String()
-								s.sweeper.removeTask(taskId)
-								log.Debugf("sweeper: unscheduled task for tx %s", taskId)
+								if err := s.sweeper.removeTask(taskId); err != nil {
+									log.WithError(err).Warnf(
+										"sweeper: failed to unschedule task for tx %s", taskId,
+									)
+								} else {
+									log.Debugf("sweeper: unscheduled task for tx %s", taskId)
+								}
 							}
 						}()
 					}

internal/core/application/sweeper.go

@@ -37,26 +37,21 @@ type sweeper struct {
 	scheduler   ports.SchedulerService
 
 	noteUriPrefix string
-
-	// cache of scheduled tasks, avoid scheduling the same sweep event multiple times
-	locker *sync.Mutex
-	// TODO move the scheduled task map to LiveStore port
-	scheduledTasks map[string]struct{}
-	ctx            context.Context
+	cache         ports.LiveStore
+	ctx           context.Context
 }
 
 func newSweeper(
 	wallet ports.WalletService, repoManager ports.RepoManager, builder ports.TxBuilder,
-	scheduler ports.SchedulerService, noteUriPrefix string,
+	scheduler ports.SchedulerService, noteUriPrefix string, cache ports.LiveStore,
 ) *sweeper {
 	return &sweeper{
 		wallet, repoManager, builder, scheduler,
-		noteUriPrefix, &sync.Mutex{}, make(map[string]struct{}), nil,
+		noteUriPrefix, cache, nil,
 	}
 }
 
 func (s *sweeper) start(ctx context.Context) error {
-	s.scheduledTasks = make(map[string]struct{})
 	s.scheduler.Start()
 
 	s.ctx = ctx
@@ -209,11 +204,11 @@ func (s *sweeper) stop() {
 	s.scheduler.Stop()
 }
 
-// removeTask update the cached map of scheduled tasks
-func (s *sweeper) removeTask(id string) {
-	s.locker.Lock()
-	defer s.locker.Unlock()
-	delete(s.scheduledTasks, id)
+// removeTask releases the claim on a scheduled task id. Callers use this
+// to cancel a pending sweep (e.g. when the underlying tx is observed
+// spent before the sweep fires).
+func (s *sweeper) removeTask(id string) error {
+	return s.cache.ScheduledTasks().Remove(s.ctx, id)
 }
 
 func (s *sweeper) scheduleCheckpointSweep(
@@ -387,29 +382,37 @@ func (s *sweeper) scheduleTask(task sweeperTask) error {
 		return task.execute()
 	}
 
-	s.locker.Lock()
-	defer s.locker.Unlock()
-
-	if _, scheduled := s.scheduledTasks[task.id]; scheduled {
+	claimed, err := s.cache.ScheduledTasks().AddIfAbsent(s.ctx, task.id)
+	if err != nil {
+		return fmt.Errorf("failed to claim scheduled task %s: %w", task.id, err)
+	}
+	if !claimed {
+		// another instance (or this one earlier) already claimed it
 		return nil
 	}
 
-	s.scheduledTasks[task.id] = struct{}{}
-
 	return s.scheduler.ScheduleTaskOnce(task.at, func() {
-		// check if the task is still scheduled before executing it
-		s.locker.Lock()
-		if _, scheduled := s.scheduledTasks[task.id]; !scheduled {
+		// Cancellation check: external code may have called removeTask
+		// because the underlying tx was already spent. If the claim is
+		// gone, skip execution.
+		has, err := s.cache.ScheduledTasks().Has(s.ctx, task.id)
+		if err != nil {
+			log.WithError(err).Errorf(
+				"sweeper: failed to check scheduled task %s, proceeding anyway", task.id,
+			)
+		} else if !has {
 			log.Debugf(
 				"sweeper: task for sweeping tx %s has been unscheduled, nothing left to do",
 				task.id,
 			)
-			s.locker.Unlock()
 			return
 		}
-		s.locker.Unlock()
 
-		s.removeTask(task.id)
+		if err := s.removeTask(task.id); err != nil {
+			log.WithError(err).Errorf(
+				"sweeper: failed to release scheduled task %s", task.id,
+			)
+		}
 
 		if err := task.execute(); err != nil {
 			log.WithError(err).Errorf("failed to execute sweep of tx %s", task.id)

internal/core/ports/live_store.go

@@ -17,6 +17,7 @@ type LiveStore interface {
 	ConfirmationSessions() ConfirmationSessionsStore
 	TreeSigingSessions() TreeSigningSessionsStore
 	BoardingInputs() BoardingInputsStore
+	ScheduledTasks() ScheduledTasksStore
 }
 
 type IntentStore interface {
@@ -89,6 +90,12 @@ type BoardingInputsStore interface {
 	DeleteSignatures(ctx context.Context, batchId string) error
 }
 
+type ScheduledTasksStore interface {
+	AddIfAbsent(ctx context.Context, id string) (bool, error)
+	Remove(ctx context.Context, id string) error
+	Has(ctx context.Context, id string) (bool, error)
+}
+
 type TimedIntent struct {
 	domain.Intent
 	BoardingInputs      []BoardingInput

internal/infrastructure/live-store/inmemory/scheduled_tasks.go

@@ -0,0 +1,43 @@
+package inmemorylivestore
+
+import (
+	"context"
+	"sync"
+
+	"github.com/arkade-os/arkd/internal/core/ports"
+)
+
+type scheduledTasksStore struct {
+	lock sync.Mutex
+	ids  map[string]struct{}
+}
+
+func NewScheduledTasksStore() ports.ScheduledTasksStore {
+	return &scheduledTasksStore{
+		ids: make(map[string]struct{}),
+	}
+}
+
+func (s *scheduledTasksStore) AddIfAbsent(_ context.Context, id string) (bool, error) {
+	s.lock.Lock()
+	defer s.lock.Unlock()
+	if _, ok := s.ids[id]; ok {
+		return false, nil
+	}
+	s.ids[id] = struct{}{}
+	return true, nil
+}
+
+func (s *scheduledTasksStore) Remove(_ context.Context, id string) error {
+	s.lock.Lock()
+	defer s.lock.Unlock()
+	delete(s.ids, id)
+	return nil
+}
+
+func (s *scheduledTasksStore) Has(_ context.Context, id string) (bool, error) {
+	s.lock.Lock()
+	defer s.lock.Unlock()
+	_, ok := s.ids[id]
+	return ok, nil
+}

internal/infrastructure/live-store/inmemory/store.go

@@ -12,6 +12,7 @@ type inMemoryLiveStore struct {
 	confirmationSessionsStore ports.ConfirmationSessionsStore
 	treeSigningSessions       ports.TreeSigningSessionsStore
 	boardingInputsStore       ports.BoardingInputsStore
+	scheduledTasksStore       ports.ScheduledTasksStore
 }
 
 func NewLiveStore(txBuilder ports.TxBuilder) ports.LiveStore {
@@ -23,6 +24,7 @@ func NewLiveStore(txBuilder ports.TxBuilder) ports.LiveStore {
 		confirmationSessionsStore: NewConfirmationSessionsStore(),
 		treeSigningSessions:       NewTreeSigningSessionsStore(),
 		boardingInputsStore:       NewBoardingInputsStore(),
+		scheduledTasksStore:       NewScheduledTasksStore(),
 	}
 }
 
@@ -47,3 +49,6 @@ func (s *inMemoryLiveStore) TreeSigingSessions() ports.TreeSigningSessionsStore
 func (s *inMemoryLiveStore) BoardingInputs() ports.BoardingInputsStore {
 	return s.boardingInputsStore
 }
+func (s *inMemoryLiveStore) ScheduledTasks() ports.ScheduledTasksStore {
+	return s.scheduledTasksStore
+}

internal/infrastructure/live-store/live_store_test.go

@@ -7,6 +7,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"sync"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -764,6 +765,78 @@ func runLiveStoreTests(t *testing.T, store ports.LiveStore) {
 		require.NoError(t, err)
 		require.Empty(t, gotSigs)
 	})
+	t.Run("ScheduledTasksStore", func(t *testing.T) {
+		ctx := t.Context()
+
+		// AddIfAbsent: first call claims the id.
+		claimed, err := store.ScheduledTasks().AddIfAbsent(ctx, "tx-abc")
+		require.NoError(t, err)
+		require.True(t, claimed, "first AddIfAbsent should succeed and return true")
+
+		has, err := store.ScheduledTasks().Has(ctx, "tx-abc")
+		require.NoError(t, err)
+		require.True(t, has)
+
+		// AddIfAbsent: second call with the same id loses the race
+		secondClaim, err := store.ScheduledTasks().AddIfAbsent(ctx, "tx-abc")
+		require.NoError(t, err)
+		require.False(t, secondClaim, "second AddIfAbsent for the same id must return false, not an error")
+
+		// Has: unknown id is false.
+		has, err = store.ScheduledTasks().Has(ctx, "never-added")
+		require.NoError(t, err)
+		require.False(t, has)
+
+		// Remove: releases the claim.
+		err = store.ScheduledTasks().Remove(ctx, "tx-abc")
+		require.NoError(t, err)
+
+		has, err = store.ScheduledTasks().Has(ctx, "tx-abc")
+		require.NoError(t, err)
+		require.False(t, has, "Has must reflect Remove")
+
+		// Remove: idempotent — unknown id is a no-op.
+		err = store.ScheduledTasks().Remove(ctx, "never-added")
+		require.NoError(t, err)
+
+		// After Remove, the same id can be re-claimed.
+		claimed, err = store.ScheduledTasks().AddIfAbsent(ctx, "tx-abc")
+		require.NoError(t, err)
+		require.True(t, claimed)
+
+		require.NoError(t, store.ScheduledTasks().Remove(ctx, "tx-abc"))
+
+		// Concurrent AddIfAbsent on the same id must produce exactly one
+		// winner. This is the load-bearing property of the whole fix:
+		// without atomicity, two arkd processes could both claim the same
+		// task and both broadcast the same sweep tx.
+		const goroutines = 100
+		var wins atomic.Int32
+		var wg sync.WaitGroup
+		start := make(chan struct{})
+
+		for range goroutines {
+			wg.Add(1)
+			go func() {
+				defer wg.Done()
+				<-start
+				claimed, err := store.ScheduledTasks().AddIfAbsent(ctx, "tx-race")
+				require.NoError(t, err)
+				if claimed {
+					wins.Add(1)
+				}
+			}()
+		}
+
+		close(start)
+		wg.Wait()
+
+		require.Equal(t, int32(1), wins.Load(),
+			"AddIfAbsent must be atomic: exactly one goroutine claims the id")
+
+		require.NoError(t, store.ScheduledTasks().Remove(ctx, "tx-race"))
+
+	})
 }
 
 type intentPushFixture struct {

internal/infrastructure/live-store/redis/round.go

@@ -14,9 +14,10 @@ import (
 )
 
 const (
-	currentRoundKey      = "currentRoundStore:round"
-	boardingInputsKey    = "boardingInputsStore:numOfInputs"
-	boardingInputSigsKey = "boardingInputsStore:signatures"
+	currentRoundKey        = "currentRoundStore:round"
+	boardingInputsKey      = "boardingInputsStore:numOfInputs"
+	boardingInputSigsKey   = "boardingInputsStore:signatures"
+	scheduledTaskKeyPrefix = "scheduledTasksStore:task"
 )
 
 type currentRoundStore struct {

internal/infrastructure/live-store/redis/scheduled_tasks.go

@@ -0,0 +1,43 @@
+package redislivestore
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/arkade-os/arkd/internal/core/ports"
+	"github.com/redis/go-redis/v9"
+)
+
+
+type scheduledTasksStore struct {
+	rdb *redis.Client
+}
+
+func NewScheduledTasksStore(rdb *redis.Client) ports.ScheduledTasksStore {
+	return &scheduledTasksStore{rdb: rdb}
+}
+
+func (s *scheduledTasksStore) AddIfAbsent(ctx context.Context, id string) (bool, error) {
+	// SETNX is atomic on the Redis server: returns true iff this call set
+	// the key. Multiple arkd processes racing to claim the same task id
+	// will see exactly one true and the rest false.
+	return s.rdb.SetNX(ctx, scheduledTaskKey(id), "1", 0).Result()
+}
+
+func (s *scheduledTasksStore) Remove(ctx context.Context, id string) error {
+	// Del returns the number of keys removed; we don't care if it was 0
+	// (Remove is idempotent per the interface contract).
+	return s.rdb.Del(ctx, scheduledTaskKey(id)).Err()
+}
+
+func (s *scheduledTasksStore) Has(ctx context.Context, id string) (bool, error) {
+	n, err := s.rdb.Exists(ctx, scheduledTaskKey(id)).Result()
+	if err != nil {
+		return false, err
+	}
+	return n > 0, nil
+}
+
+func scheduledTaskKey(id string) string {
+	return fmt.Sprintf("%s:%s", scheduledTaskKeyPrefix, id)
+}

internal/infrastructure/live-store/redis/store.go

@@ -13,6 +13,7 @@ type redisLiveStore struct {
 	confirmationSessionsStore ports.ConfirmationSessionsStore
 	treeSigningSessions       ports.TreeSigningSessionsStore
 	boardingInputsStore       ports.BoardingInputsStore
+	scheduledTasksStore       ports.ScheduledTasksStore
 }
 
 func NewLiveStore(rdb *redis.Client, builder ports.TxBuilder, numOfRetries int) ports.LiveStore {
@@ -24,6 +25,7 @@ func NewLiveStore(rdb *redis.Client, builder ports.TxBuilder, numOfRetries int)
 		confirmationSessionsStore: NewConfirmationSessionsStore(rdb, numOfRetries),
 		treeSigningSessions:       NewTreeSigningSessionsStore(rdb, numOfRetries),
 		boardingInputsStore:       NewBoardingInputsStore(rdb, numOfRetries),
+		scheduledTasksStore:       NewScheduledTasksStore(rdb),
 	}
 }
 
@@ -40,3 +42,6 @@ func (s *redisLiveStore) TreeSigingSessions() ports.TreeSigningSessionsStore {
 func (s *redisLiveStore) BoardingInputs() ports.BoardingInputsStore {
 	return s.boardingInputsStore
 }
+func (s *redisLiveStore) ScheduledTasks() ports.ScheduledTasksStore {
+	return s.scheduledTasksStore
+}