support excluding nodes when creating workloads

cluster/calcium/create.go

@@ -38,6 +38,34 @@ func (c *Calcium) CreateWorkload(ctx context.Context, opts *types.DeployOptions)
 	return c.doCreateWorkloads(ctx, opts), nil
 }
 
+// getDeployNodenames checks nodenames and excludedNodenames
+// if excludedNodenames is not set or nodenames is set, then return nodenames
+// for other cases, then node within excludedNodename list will not be returned
+func (c *Calcium) getDeployNodenames(ctx context.Context, podname string, nodenames, excludedNodenames []string) ([]string, error) {
+	if len(excludedNodenames) == 0 || len(nodenames) != 0 {
+		return nodenames, nil
+	}
+
+	allNodes, err := c.ListPodNodes(ctx, podname, map[string]string{}, false)
+	if err != nil {
+		return nil, err
+	}
+
+	excludes := map[string]struct{}{}
+	for _, n := range excludedNodenames {
+		excludes[n] = struct{}{}
+	}
+
+	rv := []string{}
+	for _, n := range allNodes {
+		if _, ok := excludes[n.Name]; ok {
+			continue
+		}
+		rv = append(rv, n.Name)
+	}
+	return rv, nil
+}
+
 // transaction: resource metadata consistency
 func (c *Calcium) doCreateWorkloads(ctx context.Context, opts *types.DeployOptions) chan *types.CreateWorkloadMessage {
 	logger := log.WithField("Calcium", "doCreateWorkloads").WithField("opts", opts)
@@ -62,12 +90,18 @@ func (c *Calcium) doCreateWorkloads(ctx context.Context, opts *types.DeployOptio
 			close(ch)
 		}()
 
+		nodenames, err := c.getDeployNodenames(ctx, opts.Podname, opts.Nodenames, opts.ExcludedNodenames)
+		if err != nil {
+			ch <- &types.CreateWorkloadMessage{Error: logger.Err(err)}
+			return
+		}
+
 		_ = utils.Txn(
 			ctx,
 
 			// if: alloc resources
 			func(ctx context.Context) error {
-				return c.withNodesLocked(ctx, opts.Podname, opts.Nodenames, opts.NodeLabels, false, func(ctx context.Context, nodeMap map[string]*types.Node) (err error) {
+				return c.withNodesLocked(ctx, opts.Podname, nodenames, opts.NodeLabels, false, func(ctx context.Context, nodeMap map[string]*types.Node) (err error) {
 					defer func() {
 						if err != nil {
 							ch <- &types.CreateWorkloadMessage{Error: logger.Err(err)}

cluster/calcium/create_test.go

@@ -373,3 +373,52 @@ func newCreateWorkloadCluster(t *testing.T) (*Calcium, []*types.Node) {
 
 	return c, nodes
 }
+
+func TestGetDeployNodenames(t *testing.T) {
+	assert := assert.New(t)
+	c := NewTestCluster()
+	store := c.store.(*storemocks.Store)
+	nodes := []*types.Node{
+		{
+			NodeMeta: types.NodeMeta{Name: "A"},
+		},
+		{
+			NodeMeta: types.NodeMeta{Name: "B"},
+		},
+		{
+			NodeMeta: types.NodeMeta{Name: "C"},
+		},
+		{
+			NodeMeta: types.NodeMeta{Name: "D"},
+		},
+	}
+
+	// error
+	store.On("GetNodesByPod", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil, errors.New("fail to list pod nodes")).Once()
+	_, err := c.getDeployNodenames(context.Background(), "pod", []string{}, []string{"A", "X"})
+	assert.Error(err)
+
+	// empty nodenames, non-empty excludeNodenames
+	store.On("GetNodesByPod", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nodes, nil).Once()
+	nodenames1, err := c.getDeployNodenames(context.Background(), "pod", []string{}, []string{"A", "B"})
+	assert.NoError(err)
+	assert.Equal([]string{"C", "D"}, nodenames1)
+
+	// empty nodenames, empty excludeNodenames
+	store.On("GetNodesByPod", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nodes, nil).Once()
+	nodenames2, err := c.getDeployNodenames(context.Background(), "pod", []string{}, []string{})
+	assert.NoError(err)
+	assert.Equal([]string{}, nodenames2)
+
+	// non-empty nodenames, empty excludeNodenames
+	store.On("GetNodesByPod", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nodes, nil).Once()
+	nodenames3, err := c.getDeployNodenames(context.Background(), "pod", []string{"O", "P"}, []string{})
+	assert.NoError(err)
+	assert.Equal([]string{"O", "P"}, nodenames3)
+
+	// non-empty nodenames
+	store.On("GetNodesByPod", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nodes, nil).Once()
+	nodenames4, err := c.getDeployNodenames(context.Background(), "pod", []string{"X", "Y"}, []string{"A", "B"})
+	assert.NoError(err)
+	assert.Equal([]string{"X", "Y"}, nodenames4)
+}