rpc.go

package main

import (
	"context"
	"errors"
	"fmt"
	"net"
	"net/rpc"
	"sync"
	"time"
)

// Goreman is RPC server
type Goreman struct {
	rpcChan chan<- *rpcMessage
}

type rpcMessage struct {
	Msg  string
	Args []string
	// sending error (if any) when the task completes
	ErrCh chan error
}

// rpcExec sends the message to the supervisor loop and waits for the result.
func (r *Goreman) rpcExec(msg string, args []string) error {
	errChan := make(chan error, 1)
	r.rpcChan <- &rpcMessage{
		Msg:   msg,
		Args:  args,
		ErrCh: errChan,
	}
	return <-errChan
}

// Start do start
func (r *Goreman) Start(args []string, ret *string) (err error) {
	defer func() {
		if r := recover(); r != nil {
			err = fmt.Errorf("%v", r)
		}
	}()
	return r.rpcExec("start", args)
}

// Stop do stop
func (r *Goreman) Stop(args []string, ret *string) (err error) {
	defer func() {
		if r := recover(); r != nil {
			err = fmt.Errorf("%v", r)
		}
	}()
	return r.rpcExec("stop", args)
}

// StopAll do stop all
func (r *Goreman) StopAll(args []string, ret *string) (err error) {
	defer func() {
		if r := recover(); r != nil {
			err = fmt.Errorf("%v", r)
		}
	}()
	for _, proc := range procs {
		if err = stopProc(proc.name, nil); err != nil {
			break
		}
	}
	return err
}

// Restart do restart
func (r *Goreman) Restart(args []string, ret *string) (err error) {
	defer func() {
		if r := recover(); r != nil {
			err = fmt.Errorf("%v", r)
		}
	}()
	return r.rpcExec("restart", args)
}

// RestartAll do restart all
func (r *Goreman) RestartAll(args []string, ret *string) (err error) {
	defer func() {
		if r := recover(); r != nil {
			err = fmt.Errorf("%v", r)
		}
	}()
	return r.rpcExec("restart", nil)
}

// List do list
func (r *Goreman) List(args []string, ret *string) (err error) {
	defer func() {
		if r := recover(); r != nil {
			err = fmt.Errorf("%v", r)
		}
	}()
	*ret = ""
	for _, proc := range procs {
		*ret += proc.name + "\n"
	}
	return err
}

// Status do status
func (r *Goreman) Status(args []string, ret *string) (err error) {
	defer func() {
		if r := recover(); r != nil {
			err = fmt.Errorf("%v", r)
		}
	}()
	mu.Lock()
	ps := make([]*procInfo, len(procs))
	copy(ps, procs)
	mu.Unlock()
	*ret = ""
	for _, proc := range ps {
		proc.mu.Lock()
		running := proc.cmd != nil
		proc.mu.Unlock()
		if running {
			*ret += "*" + proc.name + "\n"
		} else {
			*ret += " " + proc.name + "\n"
		}
	}
	return err
}

// command: run.
func run(cmd string, args []string, serverPort uint) error {
	client, err := rpc.Dial("tcp", defaultServer(serverPort))
	if err != nil {
		return err
	}
	defer client.Close()
	var ret string
	switch cmd {
	case "start":
		return client.Call("Goreman.Start", args, &ret)
	case "stop":
		return client.Call("Goreman.Stop", args, &ret)
	case "stop-all":
		return client.Call("Goreman.StopAll", args, &ret)
	case "restart":
		return client.Call("Goreman.Restart", args, &ret)
	case "restart-all":
		return client.Call("Goreman.RestartAll", args, &ret)
	case "list":
		err := client.Call("Goreman.List", args, &ret)
		fmt.Print(ret)
		return err
	case "status":
		err := client.Call("Goreman.Status", args, &ret)
		fmt.Print(ret)
		return err
	}
	return errors.New("unknown command")
}

// start rpc server.
func startServer(ctx context.Context, rpcChan chan<- *rpcMessage, listenPort uint) error {
	gm := &Goreman{
		rpcChan: rpcChan,
	}
	// use a dedicated rpc.Server so repeated startServer calls in one
	// process do not keep serving a previously registered instance.
	rpcServer := rpc.NewServer()
	if err := rpcServer.Register(gm); err != nil {
		return err
	}
	server, err := net.Listen("tcp", fmt.Sprintf("%s:%d", defaultAddr(), listenPort))
	if err != nil {
		return err
	}
	go func() {
		<-ctx.Done()
		server.Close()
	}()
	var wg sync.WaitGroup
	var acceptingConns = true
	for acceptingConns {
		conns := make(chan net.Conn, 1)
		go func() {
			conn, err := server.Accept()
			if err != nil {
				return
			}
			conns <- conn
		}()
		select {
		case <-ctx.Done():
			acceptingConns = false
			break
		case client := <-conns: // server is not canceled.
			wg.Add(1)
			go func() {
				defer wg.Done()
				rpcServer.ServeConn(client)
			}()
		}
	}
	done := make(chan struct{}, 1)
	go func() {
		wg.Wait()
		done <- struct{}{}
	}()
	select {
	case <-done:
		return nil
	case <-time.After(10 * time.Second):
		return errors.New("RPC server did not shut down in 10 seconds, quitting")
	}
}