github.com/hernad/nomad@v1.6.112/e2e/execagent/execagent.go (about)

     1  // Copyright (c) HashiCorp, Inc.
     2  // SPDX-License-Identifier: MPL-2.0
     3  
     4  package execagent
     5  
     6  import (
     7  	"fmt"
     8  	"io"
     9  	"net"
    10  	"os"
    11  	"os/exec"
    12  	"path/filepath"
    13  	"text/template"
    14  
    15  	"github.com/hernad/nomad/api"
    16  )
    17  
    18  type AgentMode int
    19  
    20  const (
    21  	// Conf enum is for configuring either a client, server, or mixed agent.
    22  	ModeClient AgentMode = 1
    23  	ModeServer AgentMode = 2
    24  	ModeBoth             = ModeClient | ModeServer
    25  )
    26  
    27  func init() {
    28  	if d := os.Getenv("NOMAD_TEST_DIR"); d != "" {
    29  		BaseDir = d
    30  	}
    31  }
    32  
    33  var (
    34  	// BaseDir is where tests will store state and can be overridden by
    35  	// setting NOMAD_TEST_DIR. Defaults to "/opt/nomadtest"
    36  	BaseDir = "/opt/nomadtest"
    37  
    38  	agentTemplate = template.Must(template.New("agent").Parse(`
    39  enable_debug = true
    40  log_level = "{{ or .LogLevel "DEBUG" }}"
    41  
    42  ports {
    43    http = {{.HTTP}}
    44    rpc  = {{.RPC}}
    45    serf = {{.Serf}}
    46  }
    47  
    48  {{ if .EnableServer }}
    49  server {
    50    enabled = true
    51    bootstrap_expect = 1
    52  }
    53  {{ end }}
    54  
    55  {{ if .EnableClient }}
    56  client {
    57    enabled = true
    58    options = {
    59      "driver.raw_exec.enable" = "1"
    60    }
    61  }
    62  {{ end }}
    63  `))
    64  )
    65  
    66  type AgentTemplateVars struct {
    67  	HTTP         int
    68  	RPC          int
    69  	Serf         int
    70  	EnableClient bool
    71  	EnableServer bool
    72  	LogLevel     string
    73  }
    74  
    75  func newAgentTemplateVars() (*AgentTemplateVars, error) {
    76  	httpPort, err := getFreePort()
    77  	if err != nil {
    78  		return nil, err
    79  	}
    80  	rpcPort, err := getFreePort()
    81  	if err != nil {
    82  		return nil, err
    83  	}
    84  	serfPort, err := getFreePort()
    85  	if err != nil {
    86  		return nil, err
    87  	}
    88  
    89  	vars := AgentTemplateVars{
    90  		HTTP: httpPort,
    91  		RPC:  rpcPort,
    92  		Serf: serfPort,
    93  	}
    94  
    95  	return &vars, nil
    96  }
    97  
    98  func writeConfig(path string, vars *AgentTemplateVars) error {
    99  	f, err := os.OpenFile(path, os.O_RDWR|os.O_CREATE, 0644)
   100  	if err != nil {
   101  		return err
   102  	}
   103  	defer f.Close()
   104  	return agentTemplate.Execute(f, vars)
   105  }
   106  
   107  // NomadAgent manages an external Nomad agent process.
   108  type NomadAgent struct {
   109  	// BinPath is the path to the Nomad binary
   110  	BinPath string
   111  
   112  	// DataDir is the path state will be saved in
   113  	DataDir string
   114  
   115  	// ConfFile is the path to the agent's conf file
   116  	ConfFile string
   117  
   118  	// Cmd is the agent process
   119  	Cmd *exec.Cmd
   120  
   121  	// Vars are the config parameters used to template
   122  	Vars *AgentTemplateVars
   123  }
   124  
   125  // NewMixedAgent creates a new Nomad agent in mixed server+client mode but does
   126  // not start the agent process until the Start() method is called.
   127  func NewMixedAgent(bin string) (*NomadAgent, error) {
   128  	if err := os.MkdirAll(BaseDir, 0755); err != nil {
   129  		return nil, err
   130  	}
   131  	dir, err := os.MkdirTemp(BaseDir, "agent")
   132  	if err != nil {
   133  		return nil, err
   134  	}
   135  
   136  	vars, err := newAgentTemplateVars()
   137  	if err != nil {
   138  		return nil, err
   139  	}
   140  	vars.EnableClient = true
   141  	vars.EnableServer = true
   142  
   143  	conf := filepath.Join(dir, "config.hcl")
   144  	if err := writeConfig(conf, vars); err != nil {
   145  		return nil, err
   146  	}
   147  
   148  	na := &NomadAgent{
   149  		BinPath:  bin,
   150  		DataDir:  dir,
   151  		ConfFile: conf,
   152  		Vars:     vars,
   153  		Cmd:      exec.Command(bin, "agent", "-config", conf, "-data-dir", dir),
   154  	}
   155  	return na, nil
   156  }
   157  
   158  // NewClientServerPair creates a pair of Nomad agents: 1 server, 1 client.
   159  func NewClientServerPair(bin string, serverOut, clientOut io.Writer) (
   160  	server *NomadAgent, client *NomadAgent, err error) {
   161  
   162  	if err := os.MkdirAll(BaseDir, 0755); err != nil {
   163  		return nil, nil, err
   164  	}
   165  
   166  	sdir, err := os.MkdirTemp(BaseDir, "server")
   167  	if err != nil {
   168  		return nil, nil, err
   169  	}
   170  
   171  	svars, err := newAgentTemplateVars()
   172  	if err != nil {
   173  		return nil, nil, err
   174  	}
   175  	svars.LogLevel = "WARN"
   176  	svars.EnableServer = true
   177  
   178  	sconf := filepath.Join(sdir, "config.hcl")
   179  	if err := writeConfig(sconf, svars); err != nil {
   180  		return nil, nil, err
   181  	}
   182  
   183  	server = &NomadAgent{
   184  		BinPath:  bin,
   185  		DataDir:  sdir,
   186  		ConfFile: sconf,
   187  		Vars:     svars,
   188  		Cmd:      exec.Command(bin, "agent", "-config", sconf, "-data-dir", sdir),
   189  	}
   190  	server.Cmd.Stdout = serverOut
   191  	server.Cmd.Stderr = serverOut
   192  
   193  	cdir, err := os.MkdirTemp(BaseDir, "client")
   194  	if err != nil {
   195  		return nil, nil, err
   196  	}
   197  
   198  	cvars, err := newAgentTemplateVars()
   199  	if err != nil {
   200  		return nil, nil, err
   201  	}
   202  	cvars.EnableClient = true
   203  
   204  	cconf := filepath.Join(cdir, "config.hcl")
   205  	if err := writeConfig(cconf, cvars); err != nil {
   206  		return nil, nil, err
   207  	}
   208  
   209  	client = &NomadAgent{
   210  		BinPath:  bin,
   211  		DataDir:  cdir,
   212  		ConfFile: cconf,
   213  		Vars:     cvars,
   214  		Cmd: exec.Command(bin, "agent",
   215  			"-config", cconf,
   216  			"-data-dir", cdir,
   217  			"-servers", fmt.Sprintf("127.0.0.1:%d", svars.RPC),
   218  		),
   219  	}
   220  	client.Cmd.Stdout = clientOut
   221  	client.Cmd.Stderr = clientOut
   222  	return
   223  }
   224  
   225  // Start the agent command.
   226  func (n *NomadAgent) Start() error {
   227  	return n.Cmd.Start()
   228  }
   229  
   230  // Stop sends an interrupt signal and returns the command's Wait error.
   231  func (n *NomadAgent) Stop() error {
   232  	if err := n.Cmd.Process.Signal(os.Interrupt); err != nil {
   233  		return err
   234  	}
   235  
   236  	return n.Cmd.Wait()
   237  }
   238  
   239  // Destroy stops the agent and removes the data dir.
   240  func (n *NomadAgent) Destroy() error {
   241  	if err := n.Stop(); err != nil {
   242  		return err
   243  	}
   244  	return os.RemoveAll(n.DataDir)
   245  }
   246  
   247  // Client returns an api.Client for the agent.
   248  func (n *NomadAgent) Client() (*api.Client, error) {
   249  	conf := api.DefaultConfig()
   250  	conf.Address = fmt.Sprintf("http://127.0.0.1:%d", n.Vars.HTTP)
   251  	return api.NewClient(conf)
   252  }
   253  
   254  func getFreePort() (int, error) {
   255  	addr, err := net.ResolveTCPAddr("tcp", "localhost:0")
   256  	if err != nil {
   257  		return 0, err
   258  	}
   259  
   260  	l, err := net.ListenTCP("tcp", addr)
   261  	if err != nil {
   262  		return 0, err
   263  	}
   264  	defer l.Close()
   265  	return l.Addr().(*net.TCPAddr).Port, nil
   266  }