worker: Support setting task types at runtime

cmd/lotus-seal-worker/info.go

@@ -2,7 +2,6 @@ package main
 
 import (
 	"fmt"
-	"github.com/filecoin-project/lotus/extern/sector-storage/sealtasks"
 	"sort"
 
 	"github.com/urfave/cli/v2"
@@ -10,6 +9,7 @@ import (
 
 	"github.com/filecoin-project/lotus/chain/types"
 	lcli "github.com/filecoin-project/lotus/cli"
+	"github.com/filecoin-project/lotus/extern/sector-storage/sealtasks"
 )
 
 var infoCmd = &cli.Command{
@@ -62,14 +62,7 @@ var infoCmd = &cli.Command{
 		fmt.Printf("Reserved memory: %s\n", types.SizeStr(types.NewInt(info.Resources.MemReserved)))
 
 		fmt.Printf("Task types: ")
-		tasks := make([]sealtasks.TaskType, 0, len(tt))
-		for taskType := range tt {
-			tasks = append(tasks, taskType)
-		}
-		sort.Slice(tasks, func(i, j int) bool {
-			return tasks[i].Less(tasks[j])
-		})
-		for _, t := range tasks {
+		for _, t := range ttList(tt) {
 			fmt.Printf("%s ", t.Short())
 		}
 		fmt.Println()
@@ -101,3 +94,14 @@ var infoCmd = &cli.Command{
 		return nil
 	},
 }
+
+func ttList(tt map[sealtasks.TaskType]struct{}) []sealtasks.TaskType {
+	tasks := make([]sealtasks.TaskType, 0, len(tt))
+	for taskType := range tt {
+		tasks = append(tasks, taskType)
+	}
+	sort.Slice(tasks, func(i, j int) bool {
+		return tasks[i].Less(tasks[j])
+	})
+	return tasks
+}

cmd/lotus-seal-worker/main.go

@@ -59,6 +59,7 @@ func main() {
 		storageCmd,
 		setCmd,
 		waitQuietCmd,
+		tasksCmd,
 	}
 
 	app := &cli.App{

cmd/lotus-seal-worker/tasks.go

@@ -0,0 +1,83 @@
+package main
+
+import (
+	"context"
+	"strings"
+
+	"github.com/urfave/cli/v2"
+	"golang.org/x/xerrors"
+
+	"github.com/filecoin-project/lotus/api"
+	lcli "github.com/filecoin-project/lotus/cli"
+	"github.com/filecoin-project/lotus/extern/sector-storage/sealtasks"
+)
+
+var tasksCmd = &cli.Command{
+	Name:  "tasks",
+	Usage: "Manage task processing",
+	Subcommands: []*cli.Command{
+		tasksEnableCmd,
+		tasksDisableCmd,
+	},
+}
+
+var allowSetting = map[sealtasks.TaskType]struct{}{
+	sealtasks.TTAddPiece:   {},
+	sealtasks.TTPreCommit1: {},
+	sealtasks.TTPreCommit2: {},
+	sealtasks.TTCommit2:    {},
+	sealtasks.TTCommit1:    {},
+	sealtasks.TTUnseal:     {},
+}
+
+var settableStr = func() string {
+	var s []string
+	for _, tt := range ttList(allowSetting) {
+		s = append(s, tt.Short())
+	}
+	return strings.Join(s, "|")
+}()
+
+var tasksEnableCmd = &cli.Command{
+	Name:      "enable",
+	Usage:     "Enable a task type",
+	ArgsUsage: "[" + settableStr + "]",
+	Action:    taskAction(api.WorkerAPI.EnableTask),
+}
+
+var tasksDisableCmd = &cli.Command{
+	Name:      "disable",
+	Usage:     "Disable a task type",
+	ArgsUsage: "[" + settableStr + "]",
+	Action:    taskAction(api.WorkerAPI.DisableTask),
+}
+
+func taskAction(tf func(a api.WorkerAPI, ctx context.Context, tt sealtasks.TaskType) error) func(cctx *cli.Context) error {
+	return func(cctx *cli.Context) error {
+		if cctx.NArg() != 0 {
+			return xerrors.Errorf("expected 1 argument")
+		}
+
+		var tt sealtasks.TaskType
+		for taskType := range allowSetting {
+			if taskType.Short() == cctx.Args().First() {
+				tt = taskType
+				break
+			}
+		}
+
+		if tt == "" {
+			return xerrors.Errorf("unknown task type '%s'", cctx.Args().First())
+		}
+
+		api, closer, err := lcli.GetWorkerAPI(cctx)
+		if err != nil {
+			return err
+		}
+		defer closer()
+
+		ctx := lcli.ReqContext(cctx)
+
+		return tf(api, ctx, tt)
+	}
+}