Allow multiple exec commands (#307)

Author: benbjohnson

cmd/litefs/config.go

@@ -21,10 +21,11 @@ import (
 
 // Config represents a configuration for the binary process.
 type Config struct {
-	Exec         string `yaml:"exec"`
-	ExitOnError  bool   `yaml:"exit-on-error"`
-	SkipSync     bool   `yaml:"skip-sync"`
-	StrictVerify bool   `yaml:"strict-verify"`
+	ExitOnError  bool `yaml:"exit-on-error"`
+	SkipSync     bool `yaml:"skip-sync"`
+	StrictVerify bool `yaml:"strict-verify"`
+
+	Exec ExecConfigSlice `yaml:"exec"`
 
 	Data    DataConfig    `yaml:"data"`
 	FUSE    FUSEConfig    `yaml:"fuse"`
@@ -56,6 +57,27 @@ func NewConfig() Config {
 	return config
 }
 
+// ExecConfigSlice represents a wrapper type for handling YAML marshaling.
+type ExecConfigSlice []*ExecConfig
+
+func (a *ExecConfigSlice) UnmarshalYAML(value *yaml.Node) error {
+	switch value.Tag {
+	case "!!str":
+		*a = ExecConfigSlice{{Cmd: value.Value}}
+		return nil
+	case "!!seq":
+		return value.Decode((*[]*ExecConfig)(a))
+	default:
+		return fmt.Errorf("invalid exec config format")
+	}
+}
+
+// ExecConfig represents a single exec command.
+type ExecConfig struct {
+	Cmd         string `yaml:"cmd"`
+	IfCandidate bool   `yaml:"if-candidate"`
+}
+
 // DataConfig represents the configuration for internal LiteFS data. This
 // includes database files as well as LTX transaction files.
 type DataConfig struct {

cmd/litefs/config_test.go

@@ -0,0 +1,79 @@
+package main_test
+
+import (
+	"bytes"
+	"strings"
+	"testing"
+
+	main "github.com/superfly/litefs/cmd/litefs"
+	"gopkg.in/yaml.v3"
+)
+
+func TestConfig(t *testing.T) {
+	t.Run("EmptyExec", func(t *testing.T) {
+		var config main.Config
+		dec := yaml.NewDecoder(strings.NewReader("exec:\n"))
+		dec.KnownFields(true)
+		if err := dec.Decode(&config); err != nil {
+			t.Fatal(err)
+		}
+
+		if got, want := len(config.Exec), 0; got != want {
+			t.Fatalf("len=%v, want %v", got, want)
+		}
+	})
+
+	t.Run("InlineExec", func(t *testing.T) {
+		var config main.Config
+		dec := yaml.NewDecoder(strings.NewReader(`exec: "run me"`))
+		dec.KnownFields(true)
+		if err := dec.Decode(&config); err != nil {
+			t.Fatal(err)
+		}
+
+		if got, want := config.Exec[0].Cmd, "run me"; got != want {
+			t.Fatalf("Cmd=%q, want %q", got, want)
+		}
+	})
+
+	t.Run("SingleExec", func(t *testing.T) {
+		buf, err := testdata.ReadFile("testdata/config/single_exec.yml")
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		var config main.Config
+		dec := yaml.NewDecoder(bytes.NewReader(buf))
+		dec.KnownFields(true)
+		if err := dec.Decode(&config); err != nil {
+			t.Fatal(err)
+		}
+
+		if got, want := config.Exec[0].Cmd, "run me"; got != want {
+			t.Fatalf("Cmd=%q, want %q", got, want)
+		}
+	})
+
+	t.Run("MultiExec", func(t *testing.T) {
+		buf, err := testdata.ReadFile("testdata/config/multi_exec.yml")
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		var config main.Config
+		dec := yaml.NewDecoder(bytes.NewReader(buf))
+		dec.KnownFields(true)
+		if err := dec.Decode(&config); err != nil {
+			t.Fatal(err)
+		}
+
+		if got, want := config.Exec[0].Cmd, "run me"; got != want {
+			t.Fatalf("Cmd=%q, want %q", got, want)
+		} else if got, want := config.Exec[0].IfCandidate, true; got != want {
+			t.Fatalf("IfCandidate=%v, want %v", got, want)
+		}
+		if got, want := config.Exec[1].Cmd, "run me too"; got != want {
+			t.Fatalf("Cmd=%q, want %q", got, want)
+		}
+	})
+}

cmd/litefs/export_test.go

@@ -1,3 +1,4 @@
+// go:build linux
 package main_test
 
 import (

cmd/litefs/import_test.go

@@ -1,3 +1,4 @@
+// go:build linux
 package main_test
 
 import (

cmd/litefs/main_test.go

@@ -1,6 +1,7 @@
 package main_test
 
 import (
+	"embed"
 	"flag"
 	"log"
 	"os"
@@ -9,6 +10,9 @@ import (
 	"github.com/superfly/litefs"
 )
 
+//go:embed testdata
+var testdata embed.FS
+
 var (
 	fuseDebug = flag.Bool("fuse.debug", false, "enable fuse debugging")
 	tracing   = flag.Bool("tracing", false, "enable trace logging")

cmd/litefs/mount_linux.go

@@ -93,7 +93,7 @@ Arguments:
 
 	// Override "exec" field if specified on the CLI.
 	if args1 != nil {
-		c.Config.Exec = strings.Join(args1, " ")
+		c.Config.Exec = ExecConfigSlice{{Cmd: strings.Join(args1, " ")}}
 	}
 
 	// Override "debug" field if specified on the CLI.
@@ -260,8 +260,11 @@ func (c *MountCommand) Run(ctx context.Context) (err error) {
 	}
 
 	// Execute subcommand, if specified in config.
-	if err := c.execCmd(ctx); err != nil {
-		return fmt.Errorf("cannot exec: %w", err)
+	// Exit if no subcommand specified.
+	if len(c.Config.Exec) > 0 {
+		if err := c.execCmds(ctx); err != nil {
+			return fmt.Errorf("cannot exec: %w", err)
+		}
 	}
 
 	return nil
@@ -384,21 +387,56 @@ func (c *MountCommand) initProxyServer(ctx context.Context) error {
 	return nil
 }
 
-func (c *MountCommand) execCmd(ctx context.Context) error {
-	// Exit if no subcommand specified.
-	if c.Config.Exec == "" {
-		return nil
+// execCmds sequentially executes the commands in the "exec" config.
+// The last command is run asynchronously and will send its exit to the execCh.
+func (c *MountCommand) execCmds(ctx context.Context) error {
+	for i, config := range c.Config.Exec {
+		args, err := shellwords.Parse(config.Cmd)
+		if err != nil {
+			return fmt.Errorf("cannot parse exec command[%d]: %w", i, err)
+		}
+		cmd, args := args[0], args[1:]
+
+		// Skip if command should only run on candidate nodes and this is a non-candidate.
+		if config.IfCandidate && !c.Store.Candidate() {
+			log.Printf("node is not a candidate, skipping command execution: %s %v", cmd, args)
+			continue
+		}
+
+		// Execute all commands synchronously except for the last one.
+		// This is to support migration commands that occur before the app start.
+		if i < len(c.Config.Exec)-1 {
+			if err := c.execSyncCmd(ctx, cmd, args); err != nil {
+				return fmt.Errorf("sync cmd: %w", err)
+			}
+		} else {
+			if err := c.execBackgroundCmd(ctx, cmd, args); err != nil {
+				return fmt.Errorf("background cmd: %w", err)
+			}
+		}
 	}
 
-	// Execute subcommand process.
-	args, err := shellwords.Parse(c.Config.Exec)
-	if err != nil {
-		return fmt.Errorf("cannot parse exec command: %w", err)
+	return nil
+}
+
+func (c *MountCommand) execSyncCmd(ctx context.Context, cmd string, args []string) error {
+	log.Printf("executing command: %s %v", cmd, args)
+
+	c.cmd = exec.CommandContext(ctx, cmd, args...)
+	c.cmd.Env = os.Environ()
+	c.cmd.Stdout = os.Stdout
+	c.cmd.Stderr = os.Stderr
+	if err := c.cmd.Run(); err != nil {
+		return fmt.Errorf("cannot run command: %w", err)
 	}
 
-	log.Printf("starting subprocess: %s %v", args[0], args[1:])
+	return nil
+}
+
+func (c *MountCommand) execBackgroundCmd(ctx context.Context, cmd string, args []string) error {
+	log.Printf("starting background subprocess: %s %v", cmd, args)
 
-	c.cmd = exec.CommandContext(ctx, args[0], args[1:]...)
+	c.cmd = exec.CommandContext(ctx, cmd, args...)
 	c.cmd.Env = os.Environ()
 	c.cmd.Stdout = os.Stdout
 	c.cmd.Stderr = os.Stderr

cmd/litefs/testdata/config/multi_exec.yml

@@ -0,0 +1,5 @@
+exec:
+  - cmd: "run me"
+    if-candidate: true
+
+  - cmd: "run me too"

cmd/litefs/testdata/config/single_exec.yml

@@ -0,0 +1,2 @@
+exec:
+  - cmd: "run me"