github.com/nycdavid/zeus@v0.0.0-20201208104106-9ba439429e03/go/processtree/slavemonitor.go (about)

     1  package processtree
     2  
     3  import (
     4  	"math/rand"
     5  	"os"
     6  	"strconv"
     7  	"syscall"
     8  
     9  	"github.com/burke/zeus/go/messages"
    10  	slog "github.com/burke/zeus/go/shinylog"
    11  	"github.com/burke/zeus/go/unixsocket"
    12  )
    13  
    14  type SlaveMonitor struct {
    15  	tree             *ProcessTree
    16  	remoteMasterFile *os.File
    17  }
    18  
    19  func Error(err string) {
    20  	// TODO
    21  	println(err)
    22  }
    23  
    24  func StartSlaveMonitor(tree *ProcessTree, fileChanges <-chan []string, done chan bool) chan bool {
    25  	quit := make(chan bool)
    26  	go func() {
    27  		localMasterFile, remoteMasterFile, err := unixsocket.Socketpair(syscall.SOCK_DGRAM)
    28  		if err != nil {
    29  			Error("Couldn't create socketpair")
    30  		}
    31  
    32  		monitor := &SlaveMonitor{tree, remoteMasterFile}
    33  		defer monitor.cleanupChildren()
    34  
    35  		localMasterSocket, err := unixsocket.NewFromFile(localMasterFile)
    36  		if err != nil {
    37  			Error("Couldn't Open UNIXSocket")
    38  		}
    39  
    40  		// We just want this unix socket to be a channel so we can select on it...
    41  		registeringFds := make(chan int, 3)
    42  		go func() {
    43  			for {
    44  				if fd, err := localMasterSocket.ReadFD(); err != nil {
    45  					slog.Error(err)
    46  				} else {
    47  					registeringFds <- fd
    48  				}
    49  			}
    50  		}()
    51  
    52  		for _, slave := range monitor.tree.SlavesByName {
    53  			go slave.Run(monitor)
    54  		}
    55  
    56  		for {
    57  			select {
    58  			case <-quit:
    59  				done <- true
    60  				return
    61  			case fd := <-registeringFds:
    62  				go monitor.slaveDidBeginRegistration(fd)
    63  			case files := <-fileChanges:
    64  				if len(files) > 0 {
    65  					tree.RestartNodesWithFeatures(files)
    66  				}
    67  			}
    68  		}
    69  	}()
    70  	return quit
    71  }
    72  
    73  func (mon *SlaveMonitor) cleanupChildren() {
    74  	for _, slave := range mon.tree.SlavesByName {
    75  		slave.ForceKill()
    76  	}
    77  }
    78  
    79  func (mon *SlaveMonitor) slaveDidBeginRegistration(fd int) {
    80  	// Having just started the process, we expect an IO, which we convert to a UNIX domain socket
    81  	fileName := strconv.Itoa(rand.Int())
    82  	slaveFile := os.NewFile(uintptr(fd), fileName)
    83  	slaveUsock, err := unixsocket.NewFromFile(slaveFile)
    84  	if err != nil {
    85  		slog.Error(err)
    86  	}
    87  
    88  	// We now expect the slave to use this fd they send us to send a Pid&Identifier Message
    89  	msg, err := slaveUsock.ReadMessage()
    90  	if err != nil {
    91  		slog.Error(err)
    92  	}
    93  	pid, parentPid, identifier, err := messages.ParsePidMessage(msg)
    94  
    95  	// And the last step before executing its action, the slave sends us a pipe it will later use to
    96  	// send us all the features it's loaded.
    97  	featurePipeFd, err := slaveUsock.ReadFD()
    98  	if err != nil {
    99  		slog.Error(err)
   100  	}
   101  
   102  	slaveNode := mon.tree.FindSlaveByName(identifier)
   103  	if slaveNode == nil {
   104  		Error("slavemonitor.go:slaveDidBeginRegistration:Unknown identifier:" + identifier)
   105  	}
   106  
   107  	slaveNode.SlaveWasInitialized(pid, parentPid, slaveUsock, featurePipeFd)
   108  }