pool.go

package gophernaut

import (
	"bufio"
	"context"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"os/exec"
	"os/signal"
	"strings"
)

// Event is basically just an enum
type Event int

// Events that can be generated by our child processes
//go:generate stringer -type=Event
const (
	Start Event = iota
	Shutdown
	PiningForTheFjords
)

func copyToLog(dst *log.Logger, src io.Reader) {
	scanner := bufio.NewScanner(src)
	for scanner.Scan() {
		dst.Print(scanner.Text())
	}
}

func startProcess(control <-chan Event, events chan<- Event, executable string) {
	procLog := log.New(os.Stdout, fmt.Sprintf("gopher-worker(%s) ", executable), log.Ldate|log.Ltime)

	commandParts := strings.Split(executable, " ")
	command := exec.Command(commandParts[0], commandParts[1:]...)
	log.Printf("Command: %v\n", command)

	stdout, err := command.StdoutPipe()
	if err != nil {
		procLog.Fatalln("Unable to connect to stdout from command...")
	}
	stderr, err := command.StderrPipe()
	if err != nil {
		procLog.Fatalln("Unable to connect to stderr from command...")
	}

	go copyToLog(procLog, stdout)
	go copyToLog(procLog, stderr)

	// Actually start the subprocess and wait for it to startup
	command.Start()

	// Push an event onto the channel for manage processes to track the new process
	events <- Start
	for {
		_, ok := <-control
		if !ok {
			fmt.Println("Killing worker process after receiving close event.")
			command.Process.Kill()
			events <- Shutdown
			break
		}
	}
}

// Worker ...
type Worker struct {
	Hostname     string
	requestCount int
	pool         *Pool
	busy         bool
}

// StartRequest marks this work as busy
func (w *Worker) StartRequest() {
	w.busy = true
	w.requestCount++
	log.Printf("Worker %s request %d starting...\n", w.Hostname, w.requestCount)
}

// CompleteRequest frees this worker
func (w *Worker) CompleteRequest() {
	w.busy = false
	w.pool.workerChannel <- w
	log.Printf("Worker %s request %d complete!\n", w.Hostname, w.requestCount)
}

// GetRequestCount accessor
func (w *Worker) GetRequestCount() int {
	return w.requestCount
}

// Pool manages the pool of Worker processes to which gophernaut dispatches
// requests.
type Pool struct {
	Executables []string
	Hostnames   []string
	Size        int

	Workers       []*Worker
	workerChannel chan *Worker

	requestCount   int
	stoppedCount   int
	processCount   int
	controlChannel chan Event
	eventsChannel  chan Event
}

// Start up the pool
func (p *Pool) Start() {
	p.controlChannel = make(chan Event)
	p.eventsChannel = make(chan Event)
	p.workerChannel = make(chan *Worker, p.Size)

	// Handle signals to try to do a graceful shutdown:
	receivedSignals := make(chan os.Signal, 1)
	signal.Notify(receivedSignals, os.Interrupt) // , syscall.SIGINT, syscall.SIGTERM)
	go func() {
		for sig := range receivedSignals {
			fmt.Printf("Received signal, %s, shutting down workers...\n", sig)
			break
		}
		close(p.controlChannel)
		signal.Stop(receivedSignals)
	}()

	log.Printf("Starting processes... %d\n", len(p.Executables))
	// Actually start some processes
	for index, executable := range p.Executables {
		log.Printf("creating worker...\n")
		w := Worker{Hostname: p.Hostnames[index], pool: p}
		log.Printf("Queuing worker...\n")
		p.workerChannel <- &w
		p.Workers = append(p.Workers, &w)
		log.Printf("Starting worker process...\n")
		go startProcess(p.controlChannel, p.eventsChannel, executable)
	}
}

// GetWorker retrieves a worker from the pool and updates the pools
// state to indicate that a request is in progress
func (p *Pool) GetWorker() *Worker {
	p.requestCount++
	log.Printf("Pool starting request %d...", p.requestCount)
	return <-p.workerChannel
}

// ManageProcesses waits for processes to start waits for graceful shutdown
func (p *Pool) ManageProcesses(s *http.Server) {
	for event := range p.eventsChannel {
		switch event {
		case Shutdown:
			p.stoppedCount++
		case Start:
			p.processCount++
		}
		if p.processCount == p.stoppedCount {
			log.Printf("%d workers stopped, shutting down.\n", p.processCount)
			c := context.Background()
			s.Shutdown(c)
		}
	}
}