better distributed lock (#307)

* refine distributed lock * adjust invoke chain to new Lock interface * add testcases Co-authored-by: Xiaofan CHEN <[email protected]>
projecteru2 · Jan 4, 2021 · 5798299 · 5798299
1 parent 8c3fc63
commit 5798299
Show file tree

Hide file tree

Showing 33 changed files with 194 additions and 111 deletions.
diff --git a/cluster/calcium/calcium_test.go b/cluster/calcium/calcium_test.go
@@ -20,9 +20,9 @@ type dummyLock struct {
 }
 
 // Lock for lock
-func (d *dummyLock) Lock(ctx context.Context) error {
+func (d *dummyLock) Lock(ctx context.Context) (context.Context, error) {
 	d.m.Lock()
-	return nil
+	return context.Background(), nil
 }
 
 // Unlock for unlock

diff --git a/cluster/calcium/capacity.go b/cluster/calcium/capacity.go
@@ -18,7 +18,7 @@ func (c *Calcium) CalculateCapacity(ctx context.Context, opts *types.DeployOptio
 		Total:          0,
 		NodeCapacities: map[string]int{},
 	}
-	return msg, c.withNodesLocked(ctx, opts.Podname, opts.Nodenames, nil, false, func(nodeMap map[string]*types.Node) error {
+	return msg, c.withNodesLocked(ctx, opts.Podname, opts.Nodenames, nil, false, func(ctx context.Context, nodeMap map[string]*types.Node) error {
 		if opts.DeployStrategy != strategy.Dummy {
 			if _, msg.NodeCapacities, err = c.doAllocResource(ctx, nodeMap, opts); err != nil {
 				return errors.WithStack(err)

diff --git a/cluster/calcium/capacity_test.go b/cluster/calcium/capacity_test.go
@@ -33,7 +33,7 @@ func TestCalculateCapacity(t *testing.T) {
 	}
 	store.On("GetNode", mock.Anything, mock.Anything).Return(node1, nil)
 	lock := &lockmocks.DistributedLock{}
-	lock.On("Lock", mock.Anything).Return(nil)
+	lock.On("Lock", mock.Anything).Return(context.TODO(), nil)
 	lock.On("Unlock", mock.Anything).Return(nil)
 	store.On("CreateLock", mock.Anything, mock.Anything).Return(lock, nil)
 	// failed by wrong resource

diff --git a/cluster/calcium/control.go b/cluster/calcium/control.go
@@ -23,7 +23,7 @@ func (c *Calcium) ControlWorkload(ctx context.Context, IDs []string, t string, f
 			go func(ID string) {
 				defer wg.Done()
 				var message []*bytes.Buffer
-				err := c.withWorkloadLocked(ctx, ID, func(workload *types.Workload) error {
+				err := c.withWorkloadLocked(ctx, ID, func(ctx context.Context, workload *types.Workload) error {
 					var err error
 					switch t {
 					case cluster.WorkloadStop:

diff --git a/cluster/calcium/control_test.go b/cluster/calcium/control_test.go
@@ -20,7 +20,7 @@ func TestControlStart(t *testing.T) {
 	ctx := context.Background()
 	store := &storemocks.Store{}
 	lock := &lockmocks.DistributedLock{}
-	lock.On("Lock", mock.Anything).Return(nil)
+	lock.On("Lock", mock.Anything).Return(context.TODO(), nil)
 	lock.On("Unlock", mock.Anything).Return(nil)
 	c.store = store
 	store.On("CreateLock", mock.Anything, mock.Anything).Return(lock, nil)
@@ -111,7 +111,7 @@ func TestControlStop(t *testing.T) {
 	ctx := context.Background()
 	store := &storemocks.Store{}
 	lock := &lockmocks.DistributedLock{}
-	lock.On("Lock", mock.Anything).Return(nil)
+	lock.On("Lock", mock.Anything).Return(context.TODO(), nil)
 	lock.On("Unlock", mock.Anything).Return(nil)
 	c.store = store
 	store.On("CreateLock", mock.Anything, mock.Anything).Return(lock, nil)
@@ -156,7 +156,7 @@ func TestControlRestart(t *testing.T) {
 	ctx := context.Background()
 	store := &storemocks.Store{}
 	lock := &lockmocks.DistributedLock{}
-	lock.On("Lock", mock.Anything).Return(nil)
+	lock.On("Lock", mock.Anything).Return(context.TODO(), nil)
 	lock.On("Unlock", mock.Anything).Return(nil)
 	c.store = store
 	store.On("CreateLock", mock.Anything, mock.Anything).Return(lock, nil)

diff --git a/cluster/calcium/copy.go b/cluster/calcium/copy.go
@@ -23,7 +23,7 @@ func (c *Calcium) Copy(ctx context.Context, opts *types.CopyOptions) (chan *type
 			wg.Add(1)
 			go func(id string, paths []string) {
 				defer wg.Done()
-				if err := c.withWorkloadLocked(ctx, id, func(workload *types.Workload) error {
+				if err := c.withWorkloadLocked(ctx, id, func(ctx context.Context, workload *types.Workload) error {
 					for _, path := range paths {
 						resp, name, err := workload.Engine.VirtualizationCopyFrom(ctx, workload.ID, path)
 						ch <- makeCopyMessage(id, name, path, err, resp)

diff --git a/cluster/calcium/copy_test.go b/cluster/calcium/copy_test.go
@@ -32,7 +32,7 @@ func TestCopy(t *testing.T) {
 	}
 	store := &storemocks.Store{}
 	lock := &lockmocks.DistributedLock{}
-	lock.On("Lock", mock.Anything).Return(nil)
+	lock.On("Lock", mock.Anything).Return(context.TODO(), nil)
 	lock.On("Unlock", mock.Anything).Return(nil)
 	store.On("CreateLock", mock.Anything, mock.Anything).Return(lock, nil)
 	c.store = store

diff --git a/cluster/calcium/create.go b/cluster/calcium/create.go
@@ -66,7 +66,7 @@ func (c *Calcium) doCreateWorkloads(ctx context.Context, opts *types.DeployOptio
 
 			// if: alloc resources
 			func(ctx context.Context) error {
-				return c.withNodesLocked(ctx, opts.Podname, opts.Nodenames, opts.NodeLabels, false, func(nodeMap map[string]*types.Node) (err error) {
+				return c.withNodesLocked(ctx, opts.Podname, opts.Nodenames, opts.NodeLabels, false, func(ctx context.Context, nodeMap map[string]*types.Node) (err error) {
 					defer func() {
 						if err != nil {
 							ch <- &types.CreateWorkloadMessage{Error: err}
@@ -102,7 +102,7 @@ func (c *Calcium) doCreateWorkloads(ctx context.Context, opts *types.DeployOptio
 			// rollback: give back resources
 			func(ctx context.Context, _ bool) (err error) {
 				for nodename, rollbackIndices := range rollbackMap {
-					if e := c.withNodeLocked(ctx, nodename, func(node *types.Node) error {
+					if e := c.withNodeLocked(ctx, nodename, func(ctx context.Context, node *types.Node) error {
 						for _, plan := range plans {
 							plan.RollbackChangesOnNode(node, rollbackIndices...) // nolint:scopelint
 						}

diff --git a/cluster/calcium/create_test.go b/cluster/calcium/create_test.go
@@ -122,13 +122,13 @@ func TestCreateWorkloadTxn(t *testing.T) {
 
 	// doAllocResource fails: MakeDeployStatus
 	lock := &lockmocks.DistributedLock{}
-	lock.On("Lock", mock.Anything).Return(nil)
+	lock.On("Lock", mock.Anything).Return(context.Background(), nil)
 	lock.On("Unlock", mock.Anything).Return(nil)
 	store.On("CreateLock", mock.Anything, mock.Anything).Return(lock, nil)
 	store.On("GetPod", mock.Anything, mock.Anything).Return(pod1, nil)
 	store.On("GetNodesByPod", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nodes, nil)
 	store.On("GetNode",
-		mock.AnythingOfType("*context.timerCtx"),
+		mock.AnythingOfType("*context.emptyCtx"),
 		mock.AnythingOfType("string"),
 	).Return(
 		func(_ context.Context, name string) (node *types.Node) {

diff --git a/cluster/calcium/dissociate.go b/cluster/calcium/dissociate.go
@@ -15,8 +15,8 @@ func (c *Calcium) DissociateWorkload(ctx context.Context, IDs []string) (chan *t
 	go func() {
 		defer close(ch)
 		for _, ID := range IDs {
-			err := c.withWorkloadLocked(ctx, ID, func(workload *types.Workload) error {
-				return c.withNodeLocked(ctx, workload.Nodename, func(node *types.Node) (err error) {
+			err := c.withWorkloadLocked(ctx, ID, func(ctx context.Context, workload *types.Workload) error {
+				return c.withNodeLocked(ctx, workload.Nodename, func(ctx context.Context, node *types.Node) (err error) {
 					return utils.Txn(
 						ctx,
 						// if

diff --git a/cluster/calcium/dissociate_test.go b/cluster/calcium/dissociate_test.go
@@ -19,7 +19,7 @@ func TestDissociateWorkload(t *testing.T) {
 	c.store = store
 
 	lock := &lockmocks.DistributedLock{}
-	lock.On("Lock", mock.Anything).Return(nil)
+	lock.On("Lock", mock.Anything).Return(context.TODO(), nil)
 	lock.On("Unlock", mock.Anything).Return(nil)
 
 	c1 := &types.Workload{

diff --git a/cluster/calcium/lock.go b/cluster/calcium/lock.go
@@ -11,12 +11,13 @@ import (
 	"github.com/projecteru2/core/types"
 )
 
-func (c *Calcium) doLock(ctx context.Context, name string, timeout time.Duration) (lock.DistributedLock, error) {
+func (c *Calcium) doLock(ctx context.Context, name string, timeout time.Duration) (lock.DistributedLock, context.Context, error) {
 	lock, err := c.store.CreateLock(name, timeout)
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
-	return lock, lock.Lock(ctx)
+	ctx, err = lock.Lock(ctx)
+	return lock, ctx, err
 }
 
 func (c *Calcium) doUnlock(ctx context.Context, lock lock.DistributedLock, msg string) error {
@@ -34,54 +35,56 @@ func (c *Calcium) doUnlockAll(ctx context.Context, locks map[string]lock.Distrib
 	}
 }
 
-func (c *Calcium) withWorkloadLocked(ctx context.Context, ID string, f func(workload *types.Workload) error) error {
-	return c.withWorkloadsLocked(ctx, []string{ID}, func(workloads map[string]*types.Workload) error {
+func (c *Calcium) withWorkloadLocked(ctx context.Context, ID string, f func(context.Context, *types.Workload) error) error {
+	return c.withWorkloadsLocked(ctx, []string{ID}, func(ctx context.Context, workloads map[string]*types.Workload) error {
 		if c, ok := workloads[ID]; ok {
-			return f(c)
+			return f(ctx, c)
 		}
 		return types.ErrWorkloadNotExists
 	})
 }
 
-func (c *Calcium) withNodeLocked(ctx context.Context, nodename string, f func(node *types.Node) error) error {
-	return c.withNodesLocked(ctx, "", []string{nodename}, nil, true, func(nodes map[string]*types.Node) error {
+func (c *Calcium) withNodeLocked(ctx context.Context, nodename string, f func(context.Context, *types.Node) error) error {
+	return c.withNodesLocked(ctx, "", []string{nodename}, nil, true, func(ctx context.Context, nodes map[string]*types.Node) error {
 		if n, ok := nodes[nodename]; ok {
-			return f(n)
+			return f(ctx, n)
 		}
 		return types.ErrNodeNotExists
 	})
 }
 
-func (c *Calcium) withWorkloadsLocked(ctx context.Context, IDs []string, f func(workloads map[string]*types.Workload) error) error {
+func (c *Calcium) withWorkloadsLocked(ctx context.Context, IDs []string, f func(context.Context, map[string]*types.Workload) error) error {
 	workloads := map[string]*types.Workload{}
 	locks := map[string]lock.DistributedLock{}
-	defer func() { c.doUnlockAll(ctx, locks) }()
+	defer func() { c.doUnlockAll(context.Background(), locks) }()
 	cs, err := c.GetWorkloads(ctx, IDs)
 	if err != nil {
 		return err
 	}
+	var lock lock.DistributedLock
 	for _, workload := range cs {
-		lock, err := c.doLock(ctx, fmt.Sprintf(cluster.WorkloadLock, workload.ID), c.config.LockTimeout)
+		lock, ctx, err = c.doLock(ctx, fmt.Sprintf(cluster.WorkloadLock, workload.ID), c.config.LockTimeout)
 		if err != nil {
 			return err
 		}
 		locks[workload.ID] = lock
 		workloads[workload.ID] = workload
 	}
-	return f(workloads)
+	return f(ctx, workloads)
 }
 
-func (c *Calcium) withNodesLocked(ctx context.Context, podname string, nodenames []string, labels map[string]string, all bool, f func(nodes map[string]*types.Node) error) error {
+func (c *Calcium) withNodesLocked(ctx context.Context, podname string, nodenames []string, labels map[string]string, all bool, f func(context.Context, map[string]*types.Node) error) error {
 	nodes := map[string]*types.Node{}
 	locks := map[string]lock.DistributedLock{}
-	defer c.doUnlockAll(ctx, locks)
+	defer c.doUnlockAll(context.Background(), locks)
 	ns, err := c.getNodes(ctx, podname, nodenames, labels, all)
 	if err != nil {
 		return err
 	}
 
+	var lock lock.DistributedLock
 	for _, n := range ns {
-		lock, err := c.doLock(ctx, fmt.Sprintf(cluster.NodeLock, podname, n.Name), c.config.LockTimeout)
+		lock, ctx, err = c.doLock(ctx, fmt.Sprintf(cluster.NodeLock, podname, n.Name), c.config.LockTimeout)
 		if err != nil {
 			return err
 		}
@@ -94,5 +97,5 @@ func (c *Calcium) withNodesLocked(ctx context.Context, podname string, nodenames
 		}
 		nodes[n.Name] = node
 	}
-	return f(nodes)
+	return f(ctx, nodes)
 }
diff --git a/cluster/calcium/lock_test.go b/cluster/calcium/lock_test.go
@@ -21,18 +21,18 @@ func TestDoLock(t *testing.T) {
 	c.store = store
 	// create lock failed
 	store.On("CreateLock", mock.Anything, mock.Anything).Return(nil, types.ErrNoETCD).Once()
-	_, err := c.doLock(ctx, "somename", 1)
+	_, _, err := c.doLock(ctx, "somename", 1)
 	assert.Error(t, err)
 
 	lock := &lockmocks.DistributedLock{}
 	store.On("CreateLock", mock.Anything, mock.Anything).Return(lock, nil)
 	// lock failed
-	lock.On("Lock", mock.Anything).Return(types.ErrNoETCD).Once()
-	_, err = c.doLock(ctx, "somename", 1)
+	lock.On("Lock", mock.Anything).Return(context.TODO(), types.ErrNoETCD).Once()
+	_, _, err = c.doLock(ctx, "somename", 1)
 	assert.Error(t, err)
 	// success
-	lock.On("Lock", mock.Anything).Return(nil)
-	_, err = c.doLock(ctx, "somename", 1)
+	lock.On("Lock", mock.Anything).Return(context.TODO(), nil)
+	_, _, err = c.doLock(ctx, "somename", 1)
 	assert.NoError(t, err)
 }
 
@@ -57,15 +57,15 @@ func TestWithWorkloadsLocked(t *testing.T) {
 	store.On("CreateLock", mock.Anything, mock.Anything).Return(lock, nil)
 	lock.On("Unlock", mock.Anything).Return(nil)
 	// failed to get lock
-	lock.On("Lock", mock.Anything).Return(types.ErrNoETCD).Once()
+	lock.On("Lock", mock.Anything).Return(context.TODO(), types.ErrNoETCD).Once()
 	store.On("GetWorkloads", mock.Anything, mock.Anything).Return([]*types.Workload{{}}, nil).Once()
-	err := c.withWorkloadsLocked(ctx, []string{"c1", "c2"}, func(workloads map[string]*types.Workload) error { return nil })
+	err := c.withWorkloadsLocked(ctx, []string{"c1", "c2"}, func(ctx context.Context, workloads map[string]*types.Workload) error { return nil })
 	assert.Error(t, err)
 	// success
-	lock.On("Lock", mock.Anything).Return(nil)
+	lock.On("Lock", mock.Anything).Return(context.TODO(), nil)
 	// failed by getworkload
 	store.On("GetWorkloads", mock.Anything, mock.Anything).Return(nil, types.ErrNoETCD).Once()
-	err = c.withWorkloadsLocked(ctx, []string{"c1", "c2"}, func(workloads map[string]*types.Workload) error { return nil })
+	err = c.withWorkloadsLocked(ctx, []string{"c1", "c2"}, func(ctx context.Context, workloads map[string]*types.Workload) error { return nil })
 	assert.Error(t, err)
 	engine := &enginemocks.API{}
 	workload := &types.Workload{
@@ -74,7 +74,7 @@ func TestWithWorkloadsLocked(t *testing.T) {
 	}
 	store.On("GetWorkloads", mock.Anything, mock.Anything).Return([]*types.Workload{workload}, nil)
 	// success
-	err = c.withWorkloadsLocked(ctx, []string{"c1", "c1"}, func(workloads map[string]*types.Workload) error {
+	err = c.withWorkloadsLocked(ctx, []string{"c1", "c1"}, func(ctx context.Context, workloads map[string]*types.Workload) error {
 		assert.Len(t, workloads, 1)
 		return nil
 	})
@@ -91,15 +91,15 @@ func TestWithWorkloadLocked(t *testing.T) {
 	store.On("CreateLock", mock.Anything, mock.Anything).Return(lock, nil)
 	lock.On("Unlock", mock.Anything).Return(nil)
 	// failed to get lock
-	lock.On("Lock", mock.Anything).Return(types.ErrNoETCD).Once()
+	lock.On("Lock", mock.Anything).Return(context.TODO(), types.ErrNoETCD).Once()
 	store.On("GetWorkloads", mock.Anything, mock.Anything).Return([]*types.Workload{{}}, nil).Once()
-	err := c.withWorkloadLocked(ctx, "c1", func(workload *types.Workload) error { return nil })
+	err := c.withWorkloadLocked(ctx, "c1", func(ctx context.Context, workload *types.Workload) error { return nil })
 	assert.Error(t, err)
 	// success
-	lock.On("Lock", mock.Anything).Return(nil)
+	lock.On("Lock", mock.Anything).Return(context.TODO(), nil)
 	// failed by getworkload
 	store.On("GetWorkloads", mock.Anything, mock.Anything).Return(nil, types.ErrNoETCD).Once()
-	err = c.withWorkloadLocked(ctx, "c1", func(workload *types.Workload) error { return nil })
+	err = c.withWorkloadLocked(ctx, "c1", func(ctx context.Context, workload *types.Workload) error { return nil })
 	assert.Error(t, err)
 	engine := &enginemocks.API{}
 	workload := &types.Workload{
@@ -108,7 +108,7 @@ func TestWithWorkloadLocked(t *testing.T) {
 	}
 	store.On("GetWorkloads", mock.Anything, mock.Anything).Return([]*types.Workload{workload}, nil)
 	// success
-	err = c.withWorkloadLocked(ctx, "c1", func(workload *types.Workload) error {
+	err = c.withWorkloadLocked(ctx, "c1", func(ctx context.Context, workload *types.Workload) error {
 		assert.Equal(t, workload.ID, "c1")
 		return nil
 	})
@@ -132,12 +132,12 @@ func TestWithNodesLocked(t *testing.T) {
 	}
 	// failed by list nodes
 	store.On("GetNodesByPod", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return([]*types.Node{}, types.ErrNoETCD).Once()
-	err := c.withNodesLocked(ctx, "test", nil, nil, false, func(nodes map[string]*types.Node) error { return nil })
+	err := c.withNodesLocked(ctx, "test", nil, nil, false, func(ctx context.Context, nodes map[string]*types.Node) error { return nil })
 	assert.Error(t, err)
 	store.On("GetNodesByPod", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return([]*types.Node{}, nil).Once()
 	// failed by filter
 	var ns map[string]*types.Node
-	err = c.withNodesLocked(ctx, "test", nil, map[string]string{"eru": "2"}, false, func(nodes map[string]*types.Node) error {
+	err = c.withNodesLocked(ctx, "test", nil, map[string]string{"eru": "2"}, false, func(ctx context.Context, nodes map[string]*types.Node) error {
 		ns = nodes
 		return nil
 	})
@@ -146,25 +146,25 @@ func TestWithNodesLocked(t *testing.T) {
 	store.On("GetNodesByPod", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return([]*types.Node{}, nil)
 	// failed by getnode
 	store.On("GetNode", mock.Anything, mock.Anything).Return(nil, types.ErrNoETCD).Once()
-	err = c.withNodesLocked(ctx, "test", []string{"test"}, nil, false, func(nodes map[string]*types.Node) error { return nil })
+	err = c.withNodesLocked(ctx, "test", []string{"test"}, nil, false, func(ctx context.Context, nodes map[string]*types.Node) error { return nil })
 	assert.Error(t, err)
 	store.On("GetNode", mock.Anything, mock.Anything).Return(node1, nil).Once()
 	// failed by lock
 	lock := &lockmocks.DistributedLock{}
 	store.On("CreateLock", mock.Anything, mock.Anything).Return(lock, nil)
 	lock.On("Unlock", mock.Anything).Return(nil)
 	// failed to get lock
-	lock.On("Lock", mock.Anything).Return(types.ErrNoETCD).Once()
-	err = c.withNodesLocked(ctx, "test", []string{"test"}, nil, false, func(nodes map[string]*types.Node) error { return nil })
+	lock.On("Lock", mock.Anything).Return(context.TODO(), types.ErrNoETCD).Once()
+	err = c.withNodesLocked(ctx, "test", []string{"test"}, nil, false, func(ctx context.Context, nodes map[string]*types.Node) error { return nil })
 	assert.Error(t, err)
-	lock.On("Lock", mock.Anything).Return(nil)
+	lock.On("Lock", mock.Anything).Return(context.TODO(), nil)
 	// failed by get locked node
 	store.On("GetNode", mock.Anything, mock.Anything).Return(nil, types.ErrNoETCD).Once()
-	err = c.withNodesLocked(ctx, "test", []string{"test"}, nil, false, func(nodes map[string]*types.Node) error { return nil })
+	err = c.withNodesLocked(ctx, "test", []string{"test"}, nil, false, func(ctx context.Context, nodes map[string]*types.Node) error { return nil })
 	assert.Error(t, err)
 	store.On("GetNode", mock.Anything, mock.Anything).Return(node1, nil)
 	// success
-	err = c.withNodesLocked(ctx, "test", []string{"test"}, nil, false, func(nodes map[string]*types.Node) error {
+	err = c.withNodesLocked(ctx, "test", []string{"test"}, nil, false, func(ctx context.Context, nodes map[string]*types.Node) error {
 		assert.Len(t, nodes, 1)
 		return nil
 	})
@@ -190,14 +190,14 @@ func TestWithNodeLocked(t *testing.T) {
 	lock := &lockmocks.DistributedLock{}
 	store.On("CreateLock", mock.Anything, mock.Anything).Return(lock, nil)
 	lock.On("Unlock", mock.Anything).Return(nil)
-	lock.On("Lock", mock.Anything).Return(nil)
+	lock.On("Lock", mock.Anything).Return(context.TODO(), nil)
 	// failed by get locked node
 	store.On("GetNode", mock.Anything, mock.Anything).Return(nil, types.ErrNoETCD).Once()
-	err := c.withNodeLocked(ctx, "test", func(node *types.Node) error { return nil })
+	err := c.withNodeLocked(ctx, "test", func(ctx context.Context, node *types.Node) error { return nil })
 	assert.Error(t, err)
 	store.On("GetNode", mock.Anything, mock.Anything).Return(node1, nil)
 	// success
-	err = c.withNodeLocked(ctx, "test", func(node *types.Node) error {
+	err = c.withNodeLocked(ctx, "test", func(ctx context.Context, node *types.Node) error {
 		assert.Equal(t, node.Name, node1.Name)
 		return nil
 	})

diff --git a/cluster/calcium/node.go b/cluster/calcium/node.go
@@ -22,7 +22,7 @@ func (c *Calcium) RemoveNode(ctx context.Context, nodename string) error {
 	if nodename == "" {
 		return types.ErrEmptyNodeName
 	}
-	return c.withNodeLocked(ctx, nodename, func(node *types.Node) error {
+	return c.withNodeLocked(ctx, nodename, func(ctx context.Context, node *types.Node) error {
 		ws, err := c.ListNodeWorkloads(ctx, node.Name, nil)
 		if err != nil {
 			return err
@@ -53,7 +53,7 @@ func (c *Calcium) SetNode(ctx context.Context, opts *types.SetNodeOptions) (*typ
 		return nil, err
 	}
 	var n *types.Node
-	return n, c.withNodeLocked(ctx, opts.Nodename, func(node *types.Node) error {
+	return n, c.withNodeLocked(ctx, opts.Nodename, func(ctx context.Context, node *types.Node) error {
 		litter.Dump(opts)
 		opts.Normalize(node)
 		n = node