feature: critical units (#17)

* feat: add "critical" and "success_codes" field to unit, #12

* feat: add SuccessCodes support to mexec

* feat: update runner signature

* feat: fexit if short unit failed with critical

* feat: fexit if long running unit failed with critical

Author: yankeguo

README.md

@@ -128,10 +128,17 @@ command:
   - once
 ```
 
-By default `once` units will block other `minit` units until finished.
+**Non-blocking**
+
+By default, `once` units will block other `minit` units until finished.
 
 Set `blocking: false` to run `once` units in background.
 
+**Critical**
+
+If `critical` field is set to `true`, `minit` will stop if this unit failed.
+
+
 ### 3.3 Type: `daemon`
 
 `daemon` units execute after `render` and `once`. It runs long-running command.
@@ -311,6 +318,31 @@ Example:
 MINIT_DISABLE=once-demo,@demo
 ```
 
+## 4.7 Critical Units
+
+If `critical` field is set to `true`, `minit` will stop if this unit failed.
+
+By specifying the `success_codes` field for `once`, `daemon` and `cron` units, `minit` will interpret exit codes within the provided list as indicative of success.
+
+Example:
+
+```yaml
+kind: once
+name: once-demo-critical
+critical: true
+command:
+  - false
+---
+kind: once
+name: once-demo-critical
+critical: true
+success_codes:
+  - 0
+  - 1
+command:
+  - false
+```
+
 ## 5. Extra Features
 
 ### 5.1 Zombie Processes Cleaning

go.mod

@@ -6,9 +6,9 @@ require (
 	github.com/robfig/cron/v3 v3.0.1
 	github.com/stretchr/testify v1.9.0
 	github.com/yankeguo/rg v1.2.0
-	golang.org/x/net v0.25.0
-	golang.org/x/sys v0.20.0
-	golang.org/x/text v0.15.0
+	golang.org/x/net v0.26.0
+	golang.org/x/sys v0.21.0
+	golang.org/x/text v0.16.0
 	gopkg.in/yaml.v3 v3.0.1
 )
 

go.sum

@@ -8,12 +8,12 @@ github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsT
 github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
 github.com/yankeguo/rg v1.2.0 h1:EhXUSK/BvO+d369Bk9CYQG2mQTWsgPKuhXHxd6KI8XM=
 github.com/yankeguo/rg v1.2.0/go.mod h1:OBQueah3CKlMksIbYKNGWJUF20pyy/mipY9NXDXlJ+c=
-golang.org/x/net v0.25.0 h1:d/OCCoBEUq33pjydKrGQhw7IlUPI2Oylr+8qLx49kac=
-golang.org/x/net v0.25.0/go.mod h1:JkAGAh7GEvH74S6FOH42FLoXpXbE/aqXSrIQjXgsiwM=
-golang.org/x/sys v0.20.0 h1:Od9JTbYCk261bKm4M/mw7AklTlFYIa0bIp9BgSm1S8Y=
-golang.org/x/sys v0.20.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
-golang.org/x/text v0.15.0 h1:h1V/4gjBv8v9cjcR6+AR5+/cIYK5N/WAgiv4xlsEtAk=
-golang.org/x/text v0.15.0/go.mod h1:18ZOQIKpY8NJVqYksKHtTdi31H5itFRjB5/qKTNYzSU=
+golang.org/x/net v0.26.0 h1:soB7SVo0PWrY4vPW/+ay0jKDNScG2X9wFeYlXIvJsOQ=
+golang.org/x/net v0.26.0/go.mod h1:5YKkiSynbBIh3p6iOc/vibscux0x38BZDkn8sCUPxHE=
+golang.org/x/sys v0.21.0 h1:rF+pYz3DAGSQAxAu1CbC7catZg4ebC4UIeIhKxBZvws=
+golang.org/x/sys v0.21.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
+golang.org/x/text v0.16.0 h1:a94ExnEXNtEwYLGJSIUxnWoxoRz/ZcCsV63ROupILh4=
+golang.org/x/text v0.16.0/go.mod h1:GhwF1Be+LQoKShO3cGOHzqOgRrGaYc9AvblQOmPVHnI=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=

main.go

@@ -166,7 +166,9 @@ func main() {
 
 	// execute short runners
 	for _, runner := range runnersS {
-		runner.Action.Do(context.Background())
+		if err = runner.Action.Do(context.Background()); err != nil {
+			return
+		}
 	}
 
 	// quick exit
@@ -178,11 +180,17 @@ func main() {
 	// run long runners
 	ctx, cancel := context.WithCancel(context.Background())
 	wg := &sync.WaitGroup{}
+	chErr := make(chan error, 1)
 
 	for _, runner := range runnersL {
 		wg.Add(1)
 		go func(runner mrunners.Runner) {
-			runner.Action.Do(ctx)
+			if err := runner.Action.Do(ctx); err != nil {
+				select {
+				case chErr <- err:
+				default:
+				}
+			}
 			wg.Done()
 		}(runner)
 	}
@@ -192,8 +200,19 @@ func main() {
 	// wait for signals
 	chSig := make(chan os.Signal, 1)
 	signal.Notify(chSig, syscall.SIGINT, syscall.SIGTERM)
-	sig := <-chSig
-	log.Printf("signal caught: %s", sig.String())
+
+	var sig os.Signal
+
+	select {
+	case sig = <-chSig:
+		log.Printf("signal caught: %s", sig.String())
+	case err = <-chErr:
+		log.Printf("critical error caught: %s", err.Error())
+	}
+
+	if sig == nil {
+		sig = syscall.SIGTERM
+	}
 
 	// shutdown context
 	cancel()

pkg/mexec/manager.go

@@ -2,6 +2,7 @@ package mexec
 
 import (
 	"errors"
+	"fmt"
 	"io"
 	"os"
 	"os/exec"
@@ -19,14 +20,14 @@ import (
 type ExecuteOptions struct {
 	Name string
 
-	Dir     string
-	Shell   string
-	Env     map[string]string
-	Command []string
-	Charset string
+	Dir          string
+	Shell        string
+	Env          map[string]string
+	Command      []string
+	Charset      string
+	SuccessCodes []int
 
-	Logger          mlog.ProcLogger
-	IgnoreExecError bool
+	Logger mlog.ProcLogger
 }
 
 type Manager interface {
@@ -35,42 +36,45 @@ type Manager interface {
 }
 
 type manager struct {
-	childPIDs    map[int]struct{}
-	childPIDLock sync.Locker
-	charsets     map[string]encoding.Encoding
+	managedPIDs    map[int]struct{}
+	managedPIDLock sync.Locker
+	charsets       map[string]encoding.Encoding
 }
 
 func NewManager() Manager {
 	return &manager{
-		childPIDs:    map[int]struct{}{},
-		childPIDLock: &sync.Mutex{},
+		managedPIDs:    map[int]struct{}{},
+		managedPIDLock: &sync.Mutex{},
 		charsets: map[string]encoding.Encoding{
 			"gb18030": simplifiedchinese.GB18030,
 			"gbk":     simplifiedchinese.GBK,
 		},
 	}
 }
 
-func (m *manager) addChildPID(fn func() (pid int, err error)) error {
-	m.childPIDLock.Lock()
-	defer m.childPIDLock.Unlock()
-	pid, err := fn()
-	if err == nil {
-		m.childPIDs[pid] = struct{}{}
+func (m *manager) StartCommand(cmd *exec.Cmd) (done func(), err error) {
+	m.managedPIDLock.Lock()
+	defer m.managedPIDLock.Unlock()
+
+	if err = cmd.Start(); err != nil {
+		return
 	}
-	return err
-}
 
-func (m *manager) delChildPID(pid int) {
-	m.childPIDLock.Lock()
-	defer m.childPIDLock.Unlock()
-	delete(m.childPIDs, pid)
+	pid := cmd.Process.Pid
+	m.managedPIDs[pid] = struct{}{}
+	done = func() {
+		m.managedPIDLock.Lock()
+		defer m.managedPIDLock.Unlock()
+		delete(m.managedPIDs, pid)
+	}
+	return
 }
 
 func (m *manager) Signal(sig os.Signal) {
-	m.childPIDLock.Lock()
-	defer m.childPIDLock.Unlock()
-	for pid := range m.childPIDs {
+	m.managedPIDLock.Lock()
+	defer m.managedPIDLock.Unlock()
+
+	for pid := range m.managedPIDs {
 		if process, _ := os.FindProcess(pid); process != nil {
 			_ = process.Signal(sig)
 		}
@@ -148,15 +152,11 @@ func (m *manager) Execute(opts ExecuteOptions) (err error) {
 	}
 
 	// start process in the same lock with signal children
-	if err = m.addChildPID(func() (pid int, err error) {
-		if err = cmd.Start(); err != nil {
-			return
-		}
-		pid = cmd.Process.Pid
-		return
-	}); err != nil {
+	var done func()
+	if done, err = m.StartCommand(cmd); err != nil {
 		return
 	}
+	defer done()
 
 	opts.Logger.Print("minit: " + opts.Name + ": process started")
 
@@ -165,17 +165,40 @@ func (m *manager) Execute(opts ExecuteOptions) (err error) {
 	go opts.Logger.Err().ReadFrom(errPipe)
 
 	// wait for process
-	if err = cmd.Wait(); err != nil {
-		opts.Logger.Error("minit: " + opts.Name + ": process exited with error: " + err.Error())
+	err = cmd.Wait()
 
-		if opts.IgnoreExecError {
-			err = nil
+	var code int
+
+	if err != nil {
+		if ee, ok := err.(*exec.ExitError); ok {
+			code = ee.ExitCode()
+		} else {
+			opts.Logger.Error("minit: " + opts.Name + ": process exited with error: " + err.Error())
+			return
 		}
-	} else {
-		opts.Logger.Print("minit: " + opts.Name + ": process exited")
 	}
 
-	m.delChildPID(cmd.Process.Pid)
+	if checkSuccessCode(opts.SuccessCodes, code) {
+		err = nil
+		opts.Logger.Print("minit: " + opts.Name + ": process exited successfully")
+		return
+	}
+
+	err = fmt.Errorf("exit code: %d is not in success_codes", code)
+
+	opts.Logger.Error("minit: " + opts.Name + ": process exited with error: " + err.Error())
 
 	return
 }
+
+func checkSuccessCode(successCodes []int, code int) bool {
+	if len(successCodes) == 0 {
+		return code == 0
+	}
+	for _, c := range successCodes {
+		if c == code {
+			return true
+		}
+	}
+	return false
+}

pkg/mexec/manager_test.go

@@ -33,8 +33,7 @@ func TestNewManager(t *testing.T) {
 		Command: []string{
 			"echo", "$AAA",
 		},
-		Logger:          logger,
-		IgnoreExecError: true,
+		Logger: logger,
 	})
 	require.NoError(t, err)
 
@@ -56,8 +55,7 @@ func TestNewManager(t *testing.T) {
 		Command: []string{
 			"sleep", "$AAA",
 		},
-		Logger:          logger,
-		IgnoreExecError: true,
+		Logger: logger,
 	})
 	require.NoError(t, err)
 

pkg/mrunners/runner.go

@@ -11,7 +11,7 @@ import (
 )
 
 type RunnerAction interface {
-	Do(ctx context.Context)
+	Do(ctx context.Context) (err error)
 }
 
 type Runner struct {

pkg/mrunners/runner_cron.go

@@ -30,30 +30,58 @@ type runnerCron struct {
 	RunnerOptions
 }
 
-func (r *runnerCron) Do(ctx context.Context) {
+func (r *runnerCron) Do(ctx context.Context) (err error) {
 	r.Print("controller started")
 	defer r.Print("controller exited")
 
 	if r.Unit.Immediate {
-		if err := r.Exec.Execute(r.Unit.ExecuteOptions(r.Logger)); err != nil {
+		if err = r.Exec.Execute(r.Unit.ExecuteOptions(r.Logger)); err != nil {
 			r.Error("failed executing: " + err.Error())
+			if r.Unit.Critical {
+				return
+			} else {
+				err = nil
+			}
 		}
 	}
 
 	cr := cron.New(cron.WithLogger(cron.PrintfLogger(r.Logger)))
-	_, err := cr.AddFunc(r.Unit.Cron, func() {
+
+	var chErr chan error
+	if r.Unit.Critical {
+		chErr = make(chan error, 1)
+	}
+
+	if _, err = cr.AddFunc(r.Unit.Cron, func() {
 		r.Print("triggered")
 		if err := r.Exec.Execute(r.Unit.ExecuteOptions(r.Logger)); err != nil {
 			r.Error("failed executing: " + err.Error())
+			if r.Unit.Critical {
+				select {
+				case chErr <- err:
+				default:
+				}
+			} else {
+				err = nil
+			}
 		}
-	})
-
-	if err != nil {
-		panic(err)
+	}); err != nil {
+		// should fail since we have checked in init
+		return
 	}
 
 	cr.Start()
 
-	<-ctx.Done()
+	if r.Unit.Critical {
+		select {
+		case <-ctx.Done():
+		case err = <-chErr:
+		}
+	} else {
+		<-ctx.Done()
+	}
+
 	<-cr.Stop().Done()
+
+	return
 }