fix interval and job cancel

management/internals/shared/grpc/token_mgr.go

@@ -165,10 +165,11 @@ func (m *TimeBasedAuthSecretsManager) runRefreshJob(job *refreshJob) {
 }
 
 func refreshInterval(ttl time.Duration) time.Duration {
-	if ttl <= 0 {
-		ttl = defaultDuration
+	interval := ttl / 4 * 3
+	if interval <= 0 {
+		interval = defaultDuration / 4 * 3
 	}
-	return ttl / 4 * 3
+	return interval
 }
 
 // SetupRefresh starts peer credentials refresh

management/internals/shared/grpc/token_refresh_scheduler.go

@@ -4,6 +4,7 @@ import (
 	"container/heap"
 	"context"
 	"sync"
+	"sync/atomic"
 	"time"
 )
 
@@ -27,6 +28,7 @@ type refreshJob struct {
 	interval  time.Duration
 	nextRun   time.Time
 	index     int
+	cancelled atomic.Bool
 }
 
 type refreshJobHeap []*refreshJob
@@ -95,6 +97,7 @@ func (s *refreshScheduler) schedule(job *refreshJob) {
 func (s *refreshScheduler) cancel(job *refreshJob) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
+	job.cancelled.Store(true)
 	if job.index >= 0 {
 		heap.Remove(&s.jobs, job.index)
 	}
@@ -147,6 +150,9 @@ func (s *refreshScheduler) loop() {
 
 func (s *refreshScheduler) worker() {
 	for job := range s.work {
+		if job.cancelled.Load() {
+			continue
+		}
 		s.run(job)
 	}
 }

management/internals/shared/grpc/token_refresh_scheduler_test.go

@@ -0,0 +1,56 @@
+package grpc
+
+import (
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestRefreshInterval(t *testing.T) {
+	defaultInterval := defaultDuration / 4 * 3
+
+	require.Equal(t, 9*time.Hour, refreshInterval(12*time.Hour))
+	require.Equal(t, defaultInterval, refreshInterval(0))
+	require.Equal(t, defaultInterval, refreshInterval(-time.Second))
+	require.Equal(t, defaultInterval, refreshInterval(3*time.Nanosecond))
+	require.Positive(t, refreshInterval(4*time.Nanosecond))
+}
+
+func TestWorkerSkipsCancelledJob(t *testing.T) {
+	var ran atomic.Int32
+	scheduler := newRefreshScheduler(func(*refreshJob) {
+		ran.Add(1)
+	})
+
+	cancelledJob := &refreshJob{interval: time.Hour}
+	cancelledJob.cancelled.Store(true)
+	liveJob := &refreshJob{interval: time.Hour}
+
+	scheduler.work <- cancelledJob
+	scheduler.work <- liveJob
+
+	require.Eventually(t, func() bool {
+		return ran.Load() == 1
+	}, 2*time.Second, 10*time.Millisecond, "live job should run exactly once, cancelled job never")
+}
+
+func TestCancelBeforeFirePreventsRun(t *testing.T) {
+	var ran atomic.Int32
+	scheduler := newRefreshScheduler(func(*refreshJob) {
+		ran.Add(1)
+	})
+
+	job := &refreshJob{interval: 50 * time.Millisecond}
+	scheduler.schedule(job)
+	scheduler.cancel(job)
+
+	time.Sleep(150 * time.Millisecond)
+	require.Zero(t, ran.Load(), "cancelled job must never fire")
+
+	scheduler.mu.Lock()
+	heapLen := len(scheduler.jobs)
+	scheduler.mu.Unlock()
+	require.Zero(t, heapLen)
+}