small pixel drawing of a pufferfish cascade

main.go

package main

import (
	"fmt"
	"log"
	"os"
	"os/signal"
	"strings"
	"syscall"
	"time"

	"git.j3s.sh/cascade/agent"
	"git.j3s.sh/cascade/ipaddr"
	"git.j3s.sh/cascade/list"
)

const usage = `cascade agent     start a cascade agent
cascade list|ls   list nodes or services
cascade members   show serf cluster members
cascade rtt       estimate latency between nodes
`

// gracefulTimeout controls how long we wait before forcefully terminating
// note that this value interacts with serf's LeavePropagateDelay config
const gracefulTimeout = 10 * time.Second

// TODO: rename agent to something cooler
func main() {
	if len(os.Args) == 1 {
		fmt.Fprintf(os.Stderr, "%s", usage)
		os.Exit(1)
	}

	run(os.Args[1], os.Args[2:])
}

func run(command string, args []string) {
	switch command {
	case "agent":
		handleAgent()
	case "list", "ls":
		list.Run(args)
	default:
		fmt.Fprintf(os.Stderr, "'%s' is not a valid command\n\n%s", command, usage)
		os.Exit(1)
	}
}

func handleAgent() {
	config := getAgentConfig()
	fmt.Printf("%+v", config)
	agent := agent.New(config)
	if err := agent.Start(); err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
	defer agent.Shutdown()
	// join any specified startup nodes
	if err := startupJoin(agent); err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
	if err := handleSignals(agent); err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
}

// handleSignals blocks until we get an exit-causing signal
func handleSignals(agent *agent.Agent) error {
	signalCh := make(chan os.Signal, 4)
	signal.Notify(signalCh, os.Interrupt, syscall.SIGTERM)

	// Wait for a signal
	var sig os.Signal
	select {
	case s := <-signalCh:
		sig = s
	case <-agent.ShutdownCh():
		// Agent is already shutdown!
		return nil
	}
	fmt.Fprintf(os.Stderr, "caught signal %s", sig)

	// Check if we should do a graceful leave
	graceful := false
	if sig == os.Interrupt || sig == syscall.SIGTERM {
		graceful = true
	}

	// Bail fast if not doing a graceful leave
	if !graceful {
		fmt.Fprintf(os.Stderr, "leave cluster with zero grace")
		return nil
	}

	// Attempt a graceful leave
	gracefulCh := make(chan struct{})
	fmt.Printf("shutting down gracefully")
	go func() {
		if err := agent.Leave(); err != nil {
			fmt.Println("Error: ", err)
			return
		}
		close(gracefulCh)
	}()

	// Wait for leave or another signal
	select {
	case <-signalCh:
		return fmt.Errorf("idfk")
	case <-time.After(gracefulTimeout):
		return fmt.Errorf("leave timed out")
	case <-gracefulCh:
		return nil
	}
}

func startupJoin(a *agent.Agent) error {
	if len(a.Config.StartJoin) == 0 {
		return nil
	}

	n, err := a.Join(a.Config.StartJoin)
	if err != nil {
		return err
	}
	if n > 0 {
		log.Printf("issue join request nodes=%d\n", n)
	}

	return nil
}

func getAgentConfig() *agent.Config {
	config := agent.DefaultConfig()
	// CASCADE_BIND=192.168.0.15:12345
	if os.Getenv("CASCADE_BIND") != "" {
		addr := ipaddr.ParseIPPort(os.Getenv("CASCADE_BIND"))
		config.BindAddr.IP = addr.IP
		if addr.Port != 0 {
			config.BindAddr.Port = addr.Port
		}
	}
	// CASCADE_JOIN=127.0.0.1,127.0.0.5
	if os.Getenv("CASCADE_JOIN") != "" {
		config.StartJoin = strings.Split(os.Getenv("CASCADE_JOIN"), ",")
	}
	// CASCADE_NAME=nostromo.j3s.sh
	if os.Getenv("CASCADE_NAME") != "" {
		config.NodeName = os.Getenv("CASCADE_NAME")
	}
	return config
}