workflow: resharding workflow: Implement checkpointing.

go/vt/vtctld/workflow.go

@@ -42,6 +42,7 @@ func initWorkflowManager(ts topo.Server) {
 
 		// Register the Horizontal Resharding workflow.
 		resharding.Register()
+
 		// Unregister the blacklisted workflows.
 		for _, name := range workflowManagerDisable {
 			workflow.Unregister(name)

go/vt/workflow/manager.go

@@ -39,6 +39,8 @@ type Factory interface {
 	// variable filled it. This Init method should fill in the
 	// Name and Data attributes, based on the provided args.
 	// This is called during the Manager.Create phase.
+	// TODO(yipeiw): We should extend the interface to pass the topology server
+	// as well. The topology server is needed in the resarding workflow.
 	Init(w *workflowpb.Workflow, args []string) error
 
 	// Instantiate loads a workflow from the proto representation

go/vt/workflow/node.go

@@ -220,6 +220,30 @@ func (n *Node) deepCopyFrom(otherNode *Node, copyChildren bool) error {
 	return nil
 }
 
+// GetChildByPath returns the child node given the relative path to this node.
+// The caller must ensure that the node tree is not modified during the call.
+func (n *Node) GetChildByPath(subPath string) (*Node, error) {
+	// Find the subnode if needed.
+	parts := strings.Split(subPath, "/")
+
+	currentNode := n
+	for i := 0; i < len(parts); i++ {
+		childPathName := parts[i]
+		found := false
+		for _, child := range currentNode.Children {
+			if child.PathName == childPathName {
+				found = true
+				currentNode = child
+				break
+			}
+		}
+		if !found {
+			return nil, fmt.Errorf("node %v has no children named %v", currentNode.Path, childPathName)
+		}
+	}
+	return currentNode, nil
+}
+
 // ActionParameters describe an action initiated by the user.
 type ActionParameters struct {
 	// Path is the path of the Node the action was performed on.
@@ -355,6 +379,7 @@ func (m *NodeManager) updateNodeAndBroadcastLocked(userNode *Node, updateChildre
 	if err != nil {
 		return err
 	}
+
 	userNode.LastChanged = time.Now().Unix()
 	if err := savedNode.deepCopyFrom(userNode, updateChildren); err != nil {
 		return err

go/vt/workflow/resharding/checkpoint.go

@@ -0,0 +1,56 @@
+package resharding
+
+import (
+	"context"
+	"sync"
+
+	"github.com/golang/protobuf/proto"
+	"github.com/youtube/vitess/go/vt/topo"
+
+	workflowpb "github.com/youtube/vitess/go/vt/proto/workflow"
+)
+
+// CheckpointWriter saves the checkpoint data into topology server.
+type CheckpointWriter struct {
+	topoServer topo.Server
+
+	// checkpointMu is used for protecting data access during checkpointing.
+	mu         sync.Mutex
+	checkpoint *workflowpb.WorkflowCheckpoint
+	wi         *topo.WorkflowInfo
+}
+
+// NewCheckpointWriter creates a CheckpointWriter.
+func NewCheckpointWriter(ts topo.Server, checkpoint *workflowpb.WorkflowCheckpoint, wi *topo.WorkflowInfo) *CheckpointWriter {
+	return &CheckpointWriter{
+		topoServer: ts,
+		checkpoint: checkpoint,
+		wi:         wi,
+	}
+}
+
+// UpdateTask updates the task status in the checkpointing copy and
+// saves the full checkpoint to the topology server.
+func (c *CheckpointWriter) UpdateTask(taskID string, status workflowpb.TaskState, err error) error {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	errorMessage := ""
+	if err != nil {
+		errorMessage = err.Error()
+	}
+
+	t := c.checkpoint.Tasks[taskID]
+	t.State = status
+	t.Error = errorMessage
+	return c.saveLocked()
+}
+
+func (c *CheckpointWriter) saveLocked() error {
+	var err error
+	c.wi.Data, err = proto.Marshal(c.checkpoint)
+	if err != nil {
+		return err
+	}
+	return c.topoServer.SaveWorkflow(context.TODO(), c.wi)
+}