github.com/anth0d/nomad@v0.0.0-20221214183521-ae3a0a2cad06/helper/subproc/subproc.go (about)

     1  package subproc
     2  
     3  import (
     4  	"bufio"
     5  	"context"
     6  	"fmt"
     7  	"io"
     8  	"os"
     9  	"time"
    10  )
    11  
    12  const (
    13  	// ExitSuccess indicates the subprocess completed successfully.
    14  	ExitSuccess = iota
    15  
    16  	// ExitFailure indicates the subprocess terminated unsuccessfully.
    17  	ExitFailure
    18  
    19  	// ExitTimeout indicates the subprocess timed out before completion.
    20  	ExitTimeout
    21  )
    22  
    23  // MainFunc is the function that runs for this sub-process.
    24  //
    25  // The return value is a process exit code.
    26  type MainFunc func() int
    27  
    28  // Do f if nomad was launched as, "nomad [name]". This process will exit without
    29  // running any other part of Nomad.
    30  func Do(name string, f MainFunc) {
    31  	if len(os.Args) > 1 && os.Args[1] == name {
    32  		os.Exit(f())
    33  	}
    34  }
    35  
    36  // Print the given message to standard error.
    37  func Print(format string, args ...any) {
    38  	_, _ = fmt.Fprintf(os.Stderr, format+"\n", args...)
    39  }
    40  
    41  // Log the given output to the logger.
    42  //
    43  // r should be a buffer containing output (typically combined stdin + stdout)
    44  // f should be an HCLogger Print method (e.g. log.Debug)
    45  func Log(r io.Reader, f func(msg string, args ...any)) {
    46  	scanner := bufio.NewScanner(r)
    47  	for scanner.Scan() {
    48  		line := scanner.Text()
    49  		f("sub-process", "OUTPUT", line)
    50  	}
    51  }
    52  
    53  // Context creates a context setup with the given timeout.
    54  func Context(timeout time.Duration) (context.Context, context.CancelFunc) {
    55  	ctx := context.Background()
    56  	return context.WithTimeout(ctx, timeout)
    57  }
    58  
    59  // SetExpiration is used to ensure the process terminates, once ctx
    60  // is complete. A short grace period is added to allow any cleanup
    61  // to happen first.
    62  func SetExpiration(ctx context.Context) {
    63  	const graceful = 5 * time.Second
    64  	go func() {
    65  		<-ctx.Done()
    66  		time.Sleep(graceful)
    67  		os.Exit(ExitTimeout)
    68  	}()
    69  }