lotus/cmd/lotus-worker/tasks.go
Phi-rjan 21d6020179
Update cmd/lotus-worker/tasks.go
Co-authored-by: Łukasz Magiera <magik6k@users.noreply.github.com>
2023-02-07 16:47:30 +01:00

116 lines
2.5 KiB
Go

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/storage/sealer/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.TTDataCid: {},
sealtasks.TTPreCommit1: {},
sealtasks.TTPreCommit2: {},
sealtasks.TTCommit2: {},
sealtasks.TTUnseal: {},
sealtasks.TTReplicaUpdate: {},
sealtasks.TTProveReplicaUpdate2: {},
sealtasks.TTRegenSectorKey: {},
}
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: "--all | [" + settableStr + "]",
Flags: []cli.Flag{
&cli.BoolFlag{
Name: "all",
Usage: "Enable all task types",
},
},
Action: taskAction(api.Worker.TaskEnable),
}
var tasksDisableCmd = &cli.Command{
Name: "disable",
Usage: "Disable a task type",
ArgsUsage: "--all | [" + settableStr + "]",
Flags: []cli.Flag{
&cli.BoolFlag{
Name: "all",
Usage: "Disable all task types",
},
},
Action: taskAction(api.Worker.TaskDisable),
}
func taskAction(tf func(a api.Worker, ctx context.Context, tt sealtasks.TaskType) error) func(cctx *cli.Context) error {
return func(cctx *cli.Context) error {
allFlag := cctx.Bool("all")
if cctx.NArg() == 1 && allFlag {
return xerrors.Errorf("Cannot use --all flag with task type argument")
}
if cctx.NArg() != 1 && !allFlag {
return xerrors.Errorf("Expected 1 argument or use --all flag")
}
var tt sealtasks.TaskType
if cctx.NArg() == 1 {
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)
if allFlag {
for taskType := range allowSetting {
if err := tf(api, ctx, taskType); err != nil {
return err
}
}
return nil
}
return tf(api, ctx, tt)
}
}