github.com/Cloud-Foundations/Dominator@v0.3.4/cmd/hypervisor/control.go (about)

     1  package main
     2  
     3  import (
     4  	"fmt"
     5  	"net"
     6  	"os"
     7  	"path/filepath"
     8  
     9  	"github.com/Cloud-Foundations/Dominator/hypervisor/manager"
    10  	"github.com/Cloud-Foundations/Dominator/lib/fsutil"
    11  	"github.com/Cloud-Foundations/Dominator/lib/log"
    12  )
    13  
    14  var shutdownVMsOnNextStop bool
    15  
    16  type flusher interface {
    17  	Flush() error
    18  }
    19  
    20  func acceptControlConnections(m *manager.Manager, listener net.Listener,
    21  	logger log.DebugLogger) {
    22  	for {
    23  		if conn, err := listener.Accept(); err != nil {
    24  			logger.Println(err)
    25  		} else if err := processControlConnection(conn, m, logger); err != nil {
    26  			logger.Println(err)
    27  		}
    28  	}
    29  }
    30  
    31  func connectToControl() net.Conn {
    32  	sockAddr := filepath.Join(*stateDir, "control")
    33  	if conn, err := net.Dial("unix", sockAddr); err != nil {
    34  		fmt.Fprintf(os.Stderr, "Error connecting to: %s: %s\n", sockAddr, err)
    35  		os.Exit(1)
    36  		return nil
    37  	} else {
    38  		return conn
    39  	}
    40  }
    41  
    42  func listenForControl(m *manager.Manager, logger log.DebugLogger) error {
    43  	sockAddr := filepath.Join(*stateDir, "control")
    44  	os.Remove(sockAddr)
    45  	if listener, err := net.Listen("unix", sockAddr); err != nil {
    46  		return err
    47  	} else {
    48  		if err := os.Chmod(sockAddr, fsutil.PrivateFilePerms); err != nil {
    49  			return err
    50  		}
    51  		go acceptControlConnections(m, listener, logger)
    52  		return nil
    53  	}
    54  }
    55  
    56  func processControlConnection(conn net.Conn, m *manager.Manager,
    57  	logger log.DebugLogger) error {
    58  	defer conn.Close()
    59  	buffer := make([]byte, 256)
    60  	if nRead, err := conn.Read(buffer); err != nil {
    61  		return fmt.Errorf("error reading request: %s", err)
    62  	} else if nRead < 1 {
    63  		return fmt.Errorf("read short request: %s", err)
    64  	} else {
    65  		request := string(buffer[:nRead])
    66  		if request[nRead-1] != '\n' {
    67  			return fmt.Errorf("request not null-terminated: %s", request)
    68  		}
    69  		request = request[:nRead-1]
    70  		switch request {
    71  		case "stop":
    72  			if _, err := fmt.Fprintln(conn, "ok"); err != nil {
    73  				return err
    74  			}
    75  			os.Remove(m.GetRootCookiePath())
    76  			if shutdownVMsOnNextStop {
    77  				logger.Println("shutting down VMs and stopping")
    78  				m.ShutdownVMsAndExit()
    79  			} else {
    80  				logger.Println("stopping without shutting down VMs")
    81  				if flusher, ok := logger.(flusher); ok {
    82  					flusher.Flush()
    83  				}
    84  				os.Exit(0)
    85  			}
    86  		case "stop-vms-on-next-stop":
    87  			if _, err := fmt.Fprintln(conn, "ok"); err != nil {
    88  				return err
    89  			}
    90  			shutdownVMsOnNextStop = true
    91  			logger.Println("will shut down VMs on next stop")
    92  		default:
    93  			if _, err := fmt.Fprintln(conn, "bad request"); err != nil {
    94  				return err
    95  			}
    96  		}
    97  	}
    98  	return nil
    99  }
   100  
   101  func sendRequest(conn net.Conn, request string) error {
   102  	if _, err := fmt.Fprintln(conn, request); err != nil {
   103  		return fmt.Errorf("error writing request: %s\n", err)
   104  	}
   105  	buffer := make([]byte, 256)
   106  	if nRead, err := conn.Read(buffer); err != nil {
   107  		return fmt.Errorf("error reading response: %s\n", err)
   108  	} else if nRead < 1 {
   109  		return fmt.Errorf("read short response: %s\n", err)
   110  	} else {
   111  		response := string(buffer[:nRead])
   112  		if response[nRead-1] != '\n' {
   113  			return fmt.Errorf("response not null-terminated: %s\n", response)
   114  		}
   115  		response = response[:nRead-1]
   116  		if response != "ok" {
   117  			return fmt.Errorf("bad response: %s\n", response)
   118  		} else {
   119  			conn.Read(buffer) // Wait for EOF.
   120  			conn.Close()
   121  			return nil
   122  		}
   123  	}
   124  }
   125  
   126  func stopSubcommand(args []string, logger log.DebugLogger) error {
   127  	return sendRequest(connectToControl(), "stop")
   128  }
   129  
   130  func stopVmsOnNextStopSubcommand(args []string, logger log.DebugLogger) error {
   131  	return sendRequest(connectToControl(), "stop-vms-on-next-stop")
   132  }