From a79303168df73cbb9c76396b539c17a2d677df3c Mon Sep 17 00:00:00 2001 From: Carl Jackson Date: Sun, 2 Mar 2014 16:49:34 -0800 Subject: [PATCH] Socket bind helper package Package bind provides a convenient syntax for binding to sockets, as well as a fair bit of magic to automatically Do The Right Thing in both development and production. --- bind/bind.go | 115 ++++++++++++++++++++++++++++++++++++++++++++++++ bind/einhorn.go | 88 ++++++++++++++++++++++++++++++++++++ bind/systemd.go | 34 ++++++++++++++ 3 files changed, 237 insertions(+) create mode 100644 bind/bind.go create mode 100644 bind/einhorn.go create mode 100644 bind/systemd.go diff --git a/bind/bind.go b/bind/bind.go new file mode 100644 index 0000000..5541ec9 --- /dev/null +++ b/bind/bind.go @@ -0,0 +1,115 @@ +/* +Package bind provides a convenient way to bind to sockets. It exposes a flag in +the default flag set named "bind" which provides syntax to bind TCP and UNIX +sockets. It also supports binding to arbitrary file descriptors passed by a +parent (for instance, systemd), and for binding to Einhorn sockets (including +Einhorn ACK support). + +If the value passed to bind contains a colon, as in ":8000" or "127.0.0.1:9001", +it will be treated as a TCP address. If it begins with a "/" or a ".", it will +be treated as a path to a UNIX socket. If it begins with the string "fd@", as in +"fd@3", it will be treated as a file descriptor (useful for use with systemd, +for instance). If it begins with the string "einhorn@", as in "einhorn@0", the +corresponding einhorn socket will be used. + +If an option is not explicitly passed, the implementation will automatically +select between using "einhorn@0", "fd@3", and ":8000", depending on whether +Einhorn or systemd (or neither) is detected. + +This package is a teensy bit magical, and goes out of its way to Do The Right +Thing in many situations, including in both development and production. If +you're looking for something less magical, you'd probably be better off just +calling net.Listen() the old-fashioned way. +*/ +package bind + +import ( + "flag" + "fmt" + "log" + "net" + "os" + "strconv" + "strings" + "sync" +) + +var bind string + +func init() { + einhornInit() + systemdInit() + + defaultBind := ":8000" + if usingEinhorn() { + defaultBind = "einhorn@0" + } else if usingSystemd() { + defaultBind = "fd@3" + } + flag.StringVar(&bind, "bind", defaultBind, + `Address to bind on. If this value has a colon, as in ":8000" or + "127.0.0.1:9001", it will be treated as a TCP address. If it + begins with a "/" or a ".", it will be treated as a path to a + UNIX socket. If it begins with the string "fd@", as in "fd@3", + it will be treated as a file descriptor (useful for use with + systemd, for instance). If it begins with the string "einhorn@", + as in "einhorn@0", the corresponding einhorn socket will be + used. If an option is not explicitly passed, the implementation + will automatically select among "einhorn@0" (Einhorn), "fd@3" + (systemd), and ":8000" (fallback) based on its environment.`) +} + +func listenTo(bind string) (net.Listener, error) { + if strings.Contains(bind, ":") { + return net.Listen("tcp", bind) + } else if strings.HasPrefix(bind, ".") || strings.HasPrefix(bind, "/") { + return net.Listen("unix", bind) + } else if strings.HasPrefix(bind, "fd@") { + fd, err := strconv.Atoi(bind[3:]) + if err != nil { + return nil, fmt.Errorf("Error while parsing fd %v: %v", + bind, err) + } + f := os.NewFile(uintptr(fd), bind) + return net.FileListener(f) + } else if strings.HasPrefix(bind, "einhorn@") { + fd, err := strconv.Atoi(bind[8:]) + if err != nil { + return nil, fmt.Errorf( + "Error while parsing einhorn %v: %v", bind, err) + } + return einhornBind(fd) + } + + return nil, fmt.Errorf("Error while parsing bind arg %v", bind) +} + +// Parse and bind to the specified address. If Socket encounters an error while +// parsing or binding to the given socket it will exit by calling log.Fatal. +func Socket(bind string) net.Listener { + l, err := listenTo(bind) + if err != nil { + log.Fatal(err) + } + return l +} + +// Parse and bind to the default socket as given to us by the flag module. If +// there was an error parsing or binding to that socket, Default will exit by +// calling `log.Fatal`. +func Default() net.Listener { + return Socket(bind) +} + +// I'm not sure why you'd ever want to call Ready() more than once, but we may +// as well be safe against it... +var ready sync.Once + +// Notify the environment (for now, just Einhorn) that the process is ready to +// receive traffic. Should be called at the last possible moment to maximize the +// chances that a faulty process exits before signaling that it's ready. +func Ready() { + ready.Do(func() { + einhornAck() + }) +} diff --git a/bind/einhorn.go b/bind/einhorn.go new file mode 100644 index 0000000..f756fb0 --- /dev/null +++ b/bind/einhorn.go @@ -0,0 +1,88 @@ +package bind + +import ( + "fmt" + "log" + "net" + "os" + "strconv" + "syscall" +) + +const tooBigErr = "bind: einhorn@%d not found (einhorn only passed %d fds)" +const bindErr = "bind: could not bind einhorn@%d: not running under einhorn" +const einhornErr = "bind: einhorn environment initialization error" +const ackErr = "bind: error ACKing to einhorn: %v" + +var einhornNumFds int + +func envInt(val string) (int, error) { + return strconv.Atoi(os.Getenv(val)) +} + +// Unfortunately this can't be a normal init function, because their execution +// order is undefined, and we need to run before the init() in bind.go. +func einhornInit() { + mpid, err := envInt("EINHORN_MASTER_PID") + if err != nil || mpid != os.Getppid() { + return + } + + einhornNumFds, err = envInt("EINHORN_FD_COUNT") + if err != nil { + einhornNumFds = 0 + return + } + + // Prevent einhorn's fds from leaking to our children + for i := 0; i < einhornNumFds; i++ { + fd := int(einhornFd(i).Fd()) + syscall.CloseOnExec(fd) + } +} + +func usingEinhorn() bool { + return einhornNumFds > 0 +} + +func einhornFd(n int) *os.File { + name := fmt.Sprintf("EINHORN_FD_%d", n) + fno, err := envInt(name) + if err != nil { + log.Fatal(einhornErr) + } + return os.NewFile(uintptr(fno), name) +} + +func einhornBind(n int) (net.Listener, error) { + if !usingEinhorn() { + return nil, fmt.Errorf(bindErr, n) + } + if n >= einhornNumFds || n < 0 { + return nil, fmt.Errorf(tooBigErr, n, einhornNumFds) + } + + f := einhornFd(n) + return net.FileListener(f) +} + +// Fun story: this is actually YAML, not JSON. +const ackMsg = `{"command":"worker:ack","pid":%d}` + "\n" + +func einhornAck() { + if !usingEinhorn() { + return + } + log.Print("bind: ACKing to einhorn") + + ctl, err := net.Dial("unix", os.Getenv("EINHORN_SOCK_PATH")) + if err != nil { + log.Fatalf(ackErr, err) + } + defer ctl.Close() + + _, err = fmt.Fprintf(ctl, ackMsg, os.Getpid()) + if err != nil { + log.Fatalf(ackErr, err) + } +} diff --git a/bind/systemd.go b/bind/systemd.go new file mode 100644 index 0000000..1648bc9 --- /dev/null +++ b/bind/systemd.go @@ -0,0 +1,34 @@ +package bind + +import ( + "os" + "syscall" +) + +const systemdMinFd = 3 + +var systemdNumFds int + +// Unfortunately this can't be a normal init function, because their execution +// order is undefined, and we need to run before the init() in bind.go. +func systemdInit() { + pid, err := envInt("LISTEN_PID") + if err != nil || pid != os.Getpid() { + return + } + + systemdNumFds, err = envInt("LISTEN_FDS") + if err != nil { + systemdNumFds = 0 + return + } + + // Prevent fds from leaking to our children + for i := 0; i < systemdNumFds; i++ { + syscall.CloseOnExec(systemdMinFd + i) + } +} + +func usingSystemd() bool { + return systemdNumFds > 0 +}