From 416167349617306700480bed13fb188f5029895f Mon Sep 17 00:00:00 2001
From: Yi Jin <yi.jin@databricks.com>
Date: Thu, 11 Apr 2024 10:19:09 -0700
Subject: [PATCH] init context in constructor

Signed-off-by: Yi Jin <yi.jin@databricks.com>
---
 pkg/pool/worker_pool.go | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/pkg/pool/worker_pool.go b/pkg/pool/worker_pool.go
index b1e39bde781..6e3e77d7ab4 100644
--- a/pkg/pool/worker_pool.go
+++ b/pkg/pool/worker_pool.go
@@ -28,26 +28,27 @@ type WorkerPool interface {
 
 type workerPool struct {
 	sync.Once
+	ctx    context.Context
 	workCh chan Work
 	cancel context.CancelFunc
 }
 
 func NewWorkerPool(workers uint) WorkerPool {
+	ctx, cancel := context.WithCancel(context.Background())
 	return &workerPool{
+		ctx:    ctx,
+		cancel: cancel,
 		workCh: make(chan Work, workers),
 	}
 }
 
 func (p *workerPool) Init() {
 	p.Do(func() {
-		ctx, cancel := context.WithCancel(context.Background())
-		p.cancel = cancel
-
 		for i := 0; i < p.Size(); i++ {
 			go func() {
 				for {
 					select {
-					case <-ctx.Done():
+					case <-p.ctx.Done():
 						// TODO: exhaust workCh before exit
 						return
 					case work := <-p.workCh: