pass in context to NewManager

ava-labs · Oct 11, 2024 · f492247 · f492247
1 parent 335b222
commit f492247
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 4 deletions.
diff --git a/vms/platformvm/validators/manager.go b/vms/platformvm/validators/manager.go
@@ -90,6 +90,7 @@ type State interface {
 }
 
 func NewManager(
+	ctx context.Context,
 	log logging.Logger,
 	cfg config.Config,
 	state State,
@@ -104,7 +105,7 @@ func NewManager(
 			TTL:     recentlyAcceptedWindowTTL,
 		},
 	)
-	go recentlyAccepted.EvictEvery(context.Background(), recentlyAcceptedEvictFrequency)
+	go recentlyAccepted.EvictEvery(ctx, recentlyAcceptedEvictFrequency)
 
 	return &manager{
 		log:              log,

diff --git a/vms/platformvm/validators/manager_benchmark_test.go b/vms/platformvm/validators/manager_benchmark_test.go
@@ -58,7 +58,11 @@ func BenchmarkGetValidatorSet(b *testing.B) {
 		Validators: vdrs,
 	})
 
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	defer cancelFunc()
+
 	m := NewManager(
+		ctx,
 		logging.NoLog{},
 		config.Config{
 			Validators: vdrs,
@@ -88,7 +92,6 @@ func BenchmarkGetValidatorSet(b *testing.B) {
 		require.NoError(addSubnetDelegator(s, subnetID, genesistest.DefaultValidatorStartTime, genesistest.DefaultValidatorEndTime, nodeIDs, currentHeight))
 	}
 
-	ctx := context.Background()
 	height, err := m.GetCurrentHeight(ctx)
 	require.NoError(err)
 	require.Equal(currentHeight, height)

diff --git a/vms/platformvm/vm.go b/vms/platformvm/vm.go
@@ -149,7 +149,9 @@ func (vm *VM) Initialize(
 		return err
 	}
 
-	validatorManager := pvalidators.NewManager(chainCtx.Log, vm.Config, vm.state, vm.metrics, &vm.clock)
+	vm.onShutdownCtx, vm.onShutdownCtxCancel = context.WithCancel(context.Background())
+
+	validatorManager := pvalidators.NewManager(vm.onShutdownCtx, chainCtx.Log, vm.Config, vm.state, vm.metrics, &vm.clock)
 	vm.State = validatorManager
 	utxoVerifier := utxo.NewVerifier(vm.ctx, &vm.clock, vm.fx)
 	vm.uptimeManager = uptime.NewManager(vm.state, &vm.clock)
@@ -199,7 +201,6 @@ func (vm *VM) Initialize(
 		return fmt.Errorf("failed to initialize network: %w", err)
 	}
 
-	vm.onShutdownCtx, vm.onShutdownCtxCancel = context.WithCancel(context.Background())
 	// TODO: Wait for this goroutine to exit during Shutdown once the platformvm
 	// has better control of the context lock.
 	go vm.Network.PushGossip(vm.onShutdownCtx)