SetNodeStatus will delete node status if ttl < 0

Author: DuodenumL

cluster/calcium/node.go

@@ -73,6 +73,12 @@ func (c *Calcium) SetNode(ctx context.Context, opts *types.SetNodeOptions) (*typ
 		n.Bypass = (opts.BypassOpt == types.TriTrue) || (opts.BypassOpt == types.TriKeep && n.Bypass)
 		if n.IsDown() {
 			logger.Errorf(ctx, "[SetNodeAvailable] node marked down: %s", opts.Nodename)
+			// remove node status
+			err := c.store.SetNodeStatus(ctx, node, -1)
+			if err != nil {
+				// don't return here
+				log.Errorf(ctx, "[SetNode] failed to set node status, err: %v", err)
+			}
 		}
 		if opts.WorkloadsDown {
 			workloads, err := c.store.ListNodeWorkloads(ctx, opts.Nodename, nil)

cluster/calcium/node_test.go

@@ -140,6 +140,7 @@ func TestSetNode(t *testing.T) {
 	store.On("CreateLock", mock.Anything, mock.Anything).Return(lock, nil)
 	lock.On("Lock", mock.Anything).Return(context.TODO(), nil)
 	lock.On("Unlock", mock.Anything).Return(nil)
+	store.On("SetNodeStatus", mock.Anything, mock.Anything, mock.Anything).Return(nil)
 
 	// fail by validating
 	_, err := c.SetNode(ctx, &types.SetNodeOptions{Nodename: ""})

store/etcdv3/node.go

@@ -290,13 +290,22 @@ func (m *Mercury) doGetNodes(ctx context.Context, kvs []*mvccpb.KeyValue, labels
 }
 
 // SetNodeStatus sets status for a node, value will expire after ttl seconds
-// ttl should be larger than 0
+// ttl < 0 means to delete node status
 // this is heartbeat of node
 func (m *Mercury) SetNodeStatus(ctx context.Context, node *types.Node, ttl int64) error {
-	if ttl <= 0 {
+	if ttl == 0 {
 		return types.ErrNodeStatusTTL
 	}
 
+	// nodenames are unique
+	statusKey := filepath.Join(nodeStatusPrefix, node.Name)
+	entityKey := fmt.Sprintf(nodeInfoKey, node.Name)
+
+	if ttl < 0 {
+		_, err := m.Delete(ctx, statusKey)
+		return err
+	}
+
 	data, err := json.Marshal(types.NodeStatus{
 		Nodename: node.Name,
 		Podname:  node.Podname,
@@ -306,9 +315,6 @@ func (m *Mercury) SetNodeStatus(ctx context.Context, node *types.Node, ttl int64
 		return err
 	}
 
-	// nodenames are unique
-	statusKey := filepath.Join(nodeStatusPrefix, node.Name)
-	entityKey := fmt.Sprintf(nodeInfoKey, node.Name)
 	return m.BindStatus(ctx, entityKey, statusKey, string(data), ttl)
 }
 

store/etcdv3/node_test.go

@@ -141,13 +141,13 @@ RdCPRPt513WozkJZZAjUSP2U
 	m.config.CertPath = "/tmp"
 	node3, err := m.doAddNode(ctx, nodename3, endpoint3, podname, ca, cert, certkey, cpu, share, memory, storage, labels, nil, nil, nil)
 	assert.NoError(t, err)
-	engine3, err := m.makeClient(ctx, node3, true)
+	engine3, err := m.makeClient(ctx, node3)
 	assert.NoError(t, err)
 	_, err = engine3.Info(ctx)
 	assert.Error(t, err)
 	// failed by get key
 	node3.Name = "nokey"
-	_, err = m.makeClient(ctx, node3, true)
+	_, err = m.makeClient(ctx, node3)
 	assert.NoError(t, err)
 }
 

store/redis/node.go

@@ -288,13 +288,21 @@ func (r *Rediaron) doGetNodes(ctx context.Context, kvs map[string]string, labels
 }
 
 // SetNodeStatus sets status for a node, value will expire after ttl seconds
-// ttl should be larger than 0
+// ttl < 0 means delete node status
 // this is heartbeat of node
 func (r *Rediaron) SetNodeStatus(ctx context.Context, node *types.Node, ttl int64) error {
-	if ttl <= 0 {
+	if ttl == 0 {
 		return types.ErrNodeStatusTTL
 	}
 
+	// nodenames are unique
+	key := filepath.Join(nodeStatusPrefix, node.Name)
+
+	if ttl < 0 {
+		_, err := r.cli.Del(ctx, key).Result()
+		return err
+	}
+
 	data, err := json.Marshal(types.NodeStatus{
 		Nodename: node.Name,
 		Podname:  node.Podname,
@@ -304,8 +312,6 @@ func (r *Rediaron) SetNodeStatus(ctx context.Context, node *types.Node, ttl int6
 		return err
 	}
 
-	// nodenames are unique
-	key := filepath.Join(nodeStatusPrefix, node.Name)
 	_, err = r.cli.Set(ctx, key, string(data), time.Duration(ttl)*time.Second).Result()
 	return err
 }

store/redis/node_test.go

@@ -138,13 +138,13 @@ RdCPRPt513WozkJZZAjUSP2U
 	s.rediaron.config.CertPath = "/tmp"
 	node3, err := s.rediaron.doAddNode(ctx, nodename3, endpoint3, podname, ca, cert, certkey, cpu, share, memory, storage, labels, nil, nil, nil)
 	s.NoError(err)
-	engine3, err := s.rediaron.makeClient(ctx, node3, true)
+	engine3, err := s.rediaron.makeClient(ctx, node3)
 	s.NoError(err)
 	_, err = engine3.Info(ctx)
 	s.Error(err)
 	// failed by get key
 	node3.Name = "nokey"
-	_, err = s.rediaron.makeClient(ctx, node3, true)
+	_, err = s.rediaron.makeClient(ctx, node3)
 	s.NoError(err)
 }