forked from cockroachdb/cockroach
-
Notifications
You must be signed in to change notification settings - Fork 0
/
sdnotify.go
155 lines (138 loc) · 3.53 KB
/
sdnotify.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
// Copyright 2016 The Cockroach Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
// implied. See the License for the specific language governing
// permissions and limitations under the License. See the AUTHORS file
// for names of contributors.
//
// Author: Ben Darnell
// Package sdnotify implements both sides of the systemd readiness
// protocol. Servers can use sdnotify.Ready() to signal that they are
// ready to receive traffic, and process managers can use
// sdnotify.Exec() to run processes that implement this protocol.
package sdnotify
import (
"fmt"
"io/ioutil"
"net"
"os"
"os/exec"
"path/filepath"
"strings"
)
const (
envName = "NOTIFY_SOCKET"
readyMsg = "READY=1"
netType = "unixgram"
)
// Ready sends a readiness signal using the systemd notification
// protocol. It should be called (once) by a server after it has
// completed its initialization (including but not necessarily limited
// to binding ports) and is ready to receive traffic.
func Ready() error {
return notifyEnv(readyMsg)
}
func notifyEnv(msg string) error {
path := os.Getenv(envName)
if path == "" {
return nil
}
return notify(path, msg)
}
func notify(path, msg string) error {
addr := net.UnixAddr{
Net: netType,
Name: path,
}
conn, err := net.DialUnix(netType, nil, &addr)
if err != nil {
return err
}
defer net.Conn(conn).Close()
_, err = conn.Write([]byte(msg))
return err
}
// Exec the given command in the background using the systemd
// notification protocol. This function returns once the command has
// either exited or signaled that it is ready. If the command exits
// with a non-zero status before signaling readiness, returns an
// exec.ExitError.
func Exec(cmd *exec.Cmd) error {
l, err := listen()
if err != nil {
return err
}
defer func() { _ = l.close() }()
if cmd.Env == nil {
// Default the environment to the parent process's, minus any
// existing versions of our variable.
varPrefix := fmt.Sprintf("%s=", envName)
for _, v := range os.Environ() {
if !strings.HasPrefix(v, varPrefix) {
cmd.Env = append(cmd.Env, v)
}
}
}
cmd.Env = append(cmd.Env, fmt.Sprintf("%s=%s", envName, l.Path))
if err := cmd.Start(); err != nil {
return err
}
// This can leak goroutines, but we don't really care because we
// always exit after calling this function.
ch := make(chan error, 2)
go func() {
ch <- l.wait()
}()
go func() {
ch <- cmd.Wait()
}()
return <-ch
}
type listener struct {
Path string
tempDir string
conn *net.UnixConn
}
func listen() (listener, error) {
dir, err := ioutil.TempDir("", "sdnotify")
if err != nil {
return listener{}, err
}
path := filepath.Join(dir, "notify.sock")
conn, err := net.ListenUnixgram(netType, &net.UnixAddr{
Net: netType,
Name: path,
})
if err != nil {
return listener{}, err
}
l := listener{
Path: path,
tempDir: dir,
conn: conn}
return l, nil
}
func (l listener) wait() error {
buf := make([]byte, len(readyMsg))
for {
n, err := l.conn.Read(buf)
if err != nil {
return err
}
if string(buf[:n]) == readyMsg {
return nil
}
}
}
func (l listener) close() error {
net.Conn(l.conn).Close()
return os.RemoveAll(l.tempDir)
}