implement internal ctx in Pool

Bisonai · Jul 10, 2024 · 48237ed · 48237ed
1 parent 5f4fabb
commit 48237ed
Show file tree

Hide file tree

Showing 4 changed files with 29 additions and 13 deletions.
diff --git a/node/pkg/fetcher/accumulator.go b/node/pkg/fetcher/accumulator.go
@@ -33,10 +33,14 @@ func (a *Accumulator) Run(ctx context.Context) {
 	for {
 		select {
 		case <-ticker.C:
-			p.AddJob(ctx, func() {
+			p.AddJob(func() {
 				a.accumulatorJob(accumulatorCtx)
 			})
 		case <-ctx.Done():
+			if p.IsRunning {
+				p.Cancel()
+				p.IsRunning = false
+			}
 			log.Debug().Str("Player", "Fetcher").Msg("fetcher local aggregates channel goroutine stopped")
 			return
 		}

diff --git a/node/pkg/raft/raft.go b/node/pkg/raft/raft.go
@@ -331,6 +331,10 @@ func (r *Raft) becomeLeader(ctx context.Context) {
 				log.Debug().Msg("resigning as leader")
 				r.HeartbeatTicker.Stop()
 				r.LeaderJobTicker.Stop()
+				if p.IsRunning {
+					p.Cancel()
+					p.IsRunning = false
+				}
 
 				return
 
@@ -341,7 +345,7 @@ func (r *Raft) becomeLeader(ctx context.Context) {
 				}
 
 			case <-r.LeaderJobTicker.C:
-				p.AddJob(ctx, func() {
+				p.AddJob(func() {
 					defer func() {
 						if r := recover(); r != nil {
 							log.Error().Msgf("recovered from panic in leader job: %v", r)

diff --git a/node/pkg/utils/pool/pool.go b/node/pkg/utils/pool/pool.go
@@ -7,6 +7,9 @@ import (
 type Pool struct {
 	jobChannel  chan func()
 	workerCount int
+	ctx         context.Context
+	Cancel      context.CancelFunc
+	IsRunning   bool
 }
 
 func NewPool(workerCount int) *Pool {
@@ -17,27 +20,32 @@ func NewPool(workerCount int) *Pool {
 }
 
 func (p *Pool) Run(ctx context.Context) {
+	poolCtx, cancel := context.WithCancel(ctx)
+	p.ctx = poolCtx
+	p.Cancel = cancel
+	p.IsRunning = true
+
 	for i := 0; i < p.workerCount; i++ {
-		go p.worker(ctx)
+		go p.worker()
 	}
 }
 
-func (p *Pool) worker(ctx context.Context) {
+func (p *Pool) worker() {
 	for {
 		select {
 		case job := <-p.jobChannel:
 			job()
-		case <-ctx.Done():
+		case <-p.ctx.Done():
 			return
 		}
 	}
 }
 
-func (p *Pool) AddJob(ctx context.Context, job func()) {
+func (p *Pool) AddJob(job func()) {
 	select {
 	case p.jobChannel <- job:
 		return
-	case <-ctx.Done():
+	case <-p.ctx.Done():
 		return
 	}
 }
diff --git a/node/pkg/utils/tests/pool_test.go b/node/pkg/utils/tests/pool_test.go
@@ -25,7 +25,7 @@ func TestImmediateJobExecution(t *testing.T) {
 	p.Run(ctx)
 
 	done := make(chan bool)
-	p.AddJob(ctx, func() {
+	p.AddJob(func() {
 		done <- true
 	})
 
@@ -48,7 +48,7 @@ func TestLargeNumberOfJobs(t *testing.T) {
 	wg.Add(jobCount)
 
 	for i := 0; i < jobCount; i++ {
-		p.AddJob(ctx, func() {
+		p.AddJob(func() {
 			wg.Done()
 		})
 	}
@@ -74,7 +74,7 @@ func TestContextCancelDuringJobExecution(t *testing.T) {
 	p.Run(ctx)
 
 	done := make(chan bool)
-	p.AddJob(ctx, func() {
+	p.AddJob(func() {
 		time.Sleep(200 * time.Millisecond)
 		done <- true
 	})
@@ -104,7 +104,7 @@ func TestAddJobToClosedPool(t *testing.T) {
 
 	done := make(chan bool)
 	go func() {
-		p.AddJob(ctx, func() {
+		p.AddJob(func() {
 			done <- true
 		})
 	}()
@@ -129,7 +129,7 @@ func TestConcurrentJobExecution(t *testing.T) {
 	wg.Add(jobCount)
 
 	for i := 0; i < jobCount; i++ {
-		p.AddJob(ctx, func() {
+		p.AddJob(func() {
 			channel <- i
 			wg.Done()
 		})
@@ -163,7 +163,7 @@ func TestWorkerCount(t *testing.T) {
 
 	done := make(chan bool)
 	for i := 0; i < POOL_WORKER_COUNT; i++ {
-		p.AddJob(ctx, func() {
+		p.AddJob(func() {
 			mu.Lock()
 			workerCount++
 			mu.Unlock()