phpspyd.go

package main

import (
	"flag"
	"os"
	"os/exec"
	"os/signal"
	"strconv"
	"strings"
	"syscall"
	"time"
)

type Config struct {
	maxProcs   int
	pgrepArgs  string
	phpspyArgs string
}

var (
	config        Config
	newPidChan    chan int
	deadPidChan   chan int
	workerFinChan chan int
	doneChan      chan int
	curPidMap     map[int]bool
)

// Program entry point
func main() {
	// Parse flags
	flag.IntVar(&config.maxProcs, "maxProcs", 64, "Max phpspy procs to spawn")
	flag.StringVar(&config.pgrepArgs, "pgrepArgs", "-x httpd", "Arguments to pass to pgrep for finding pids to attach to")
	flag.StringVar(&config.phpspyArgs, "phpspyArgs", "", `Arguments to pass to phpspy in addition to "-p <pid>"`)
	flag.Parse()

	// Make stuff
	newPidChan = make(chan int, config.maxProcs)
	deadPidChan = make(chan int, config.maxProcs)
	workerFinChan = make(chan int)
	doneChan = make(chan int)
	curPidMap = make(map[int]bool)

	// Handle SIGINT and SIGTERM
	handleSignals()

	// Spawn N threads
	// Stuff a bogus pid in `deadPidChan` to trigger `getNewPid` below
	for i := 0; i < config.maxProcs; i++ {
		deadPidChan <- 0
		go work()
	}

loop:
	// Loop until doneChan is closed and N workerFinChan tokens are consumed
	for {
		select {
		case deadPid := <-deadPidChan:
			// A worker is finished with `deadPid`
			// Delete it from `curPidMap` and feed a new one into `newPidChan`
			delete(curPidMap, deadPid)
			if newPid := getNewPid(curPidMap); newPid > 0 {
				curPidMap[newPid] = true
				newPidChan <- newPid
			}
		case <-doneChan:
			// Someone sent us a SIGTERM or SIGINT
			// Close `newPidChan` and wait for N `workerFinChan` tokens before breaking
			close(newPidChan)
			for i := 0; i < config.maxProcs; i++ {
				<-workerFinChan
			}
			break loop
		}
	}
}

// Attach phpspy to pids from `newPidChan` in a loop
// Run until `newPidChan` is closed
func work() {
	// Loop over `newPidChan` until it is closed
	for pid := range newPidChan {
		// Run phpspy on `pid`
		args := strings.Split(config.phpspyArgs, " ")
		args = append(args, "-p", strconv.Itoa(pid))
		cmd := exec.Command("phpspy", args...)
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr
		// Wait for command to finish or `doneChan`
		cmdChan := make(chan int)
		go func() {
			cmd.Run()
			cmdChan <- 1
		}()
		select {
		case <-cmdChan:
		case <-doneChan:
			// Send phpspy a kill signal
			cmd.Process.Kill()
		}
		// Signal to main thread that we are done with `pid` and need a new one
		deadPidChan <- pid
	}
	// Signal to main thread that we are completely done
	workerFinChan <- 1
}

// Return a pid to attach to
// Return 0 if `doneChan` is closed
func getNewPid(curPidMap map[int]bool) int {
	// Use a closure to retain state of `newPids`
	var newPids []int
	return func() int {
		var newPid int
		for len(newPids) < 1 {
			// `newPids` is empty so pgrep for new pids
			newPids = pgrepForNewPids(curPidMap)
			if len(newPids) < 1 {
				// Still nothing, sleep a bit and try again
				time.Sleep(1 * time.Second)
				// Exit if `doneChan` is closed
				select {
				case _, ok := <-doneChan:
					if !ok {
						return 0
					}
				default:
				}
			}
		}
		// Dequeue a pid from `newPids` and return
		newPid, newPids = newPids[0], newPids[1:]
		return newPid
	}()
}

// Return a list of `pgrepForAllPids` pids excluding ones already in `curPidMap`
func pgrepForNewPids(curPidMap map[int]bool) []int {
	intPids := []int{}
	for _, strPid := range pgrepForAllPids() {
		intPid, _ := strconv.Atoi(strPid)
		if intPid != 0 && !curPidMap[intPid] {
			intPids = append(intPids, intPid)
		}
	}
	return intPids
}

// Return output from `pgrep`
func pgrepForAllPids() []string {
	cmd := exec.Command("sh", "-c", "pgrep "+config.pgrepArgs)
	cmd.Stderr = os.Stderr
	out, err := cmd.Output()
	if err != nil {
		return []string{}
	}
	return strings.Split(string(out[:]), "\n")
}

// Close `doneChan` on SIGINT or SIGTERM
func handleSignals() {
	c := make(chan os.Signal, 1)
	signal.Notify(c, syscall.SIGINT, syscall.SIGTERM)
	go func() {
		<-c
		close(doneChan)
	}()
}