github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/util/sdnotify/sdnotify_unix.go (about)

     1  // Copyright 2016 The Cockroach Authors.
     2  //
     3  // Use of this software is governed by the Business Source License
     4  // included in the file licenses/BSL.txt.
     5  //
     6  // As of the Change Date specified in that file, in accordance with
     7  // the Business Source License, use of this software will be governed
     8  // by the Apache License, Version 2.0, included in the file
     9  // licenses/APL.txt.
    10  
    11  // +build !windows
    12  
    13  package sdnotify
    14  
    15  import (
    16  	"fmt"
    17  	"io/ioutil"
    18  	"net"
    19  	"os"
    20  	"os/exec"
    21  	"path/filepath"
    22  	"strings"
    23  )
    24  
    25  const (
    26  	envName  = "NOTIFY_SOCKET"
    27  	readyMsg = "READY=1"
    28  	netType  = "unixgram"
    29  )
    30  
    31  func ready() error {
    32  	return notifyEnv(readyMsg)
    33  }
    34  
    35  func notifyEnv(msg string) error {
    36  	if path, ok := os.LookupEnv(envName); ok {
    37  		return notify(path, msg)
    38  	}
    39  	return nil
    40  }
    41  
    42  func notify(path, msg string) error {
    43  	addr := net.UnixAddr{
    44  		Net:  netType,
    45  		Name: path,
    46  	}
    47  	conn, err := net.DialUnix(netType, nil, &addr)
    48  	if err != nil {
    49  		return err
    50  	}
    51  	defer net.Conn(conn).Close()
    52  
    53  	_, err = conn.Write([]byte(msg))
    54  	return err
    55  }
    56  
    57  func bgExec(cmd *exec.Cmd) error {
    58  	l, err := listen()
    59  	if err != nil {
    60  		return err
    61  	}
    62  	defer func() { _ = l.close() }()
    63  
    64  	if cmd.Env == nil {
    65  		// Default the environment to the parent process's, minus any
    66  		// existing versions of our variable.
    67  		varPrefix := fmt.Sprintf("%s=", envName)
    68  		for _, v := range os.Environ() {
    69  			if !strings.HasPrefix(v, varPrefix) {
    70  				cmd.Env = append(cmd.Env, v)
    71  			}
    72  		}
    73  	}
    74  	cmd.Env = append(cmd.Env, fmt.Sprintf("%s=%s", envName, l.Path))
    75  
    76  	if err := cmd.Start(); err != nil {
    77  		return err
    78  	}
    79  
    80  	// This can leak goroutines, but we don't really care because we
    81  	// always exit after calling this function.
    82  	ch := make(chan error, 2)
    83  	go func() {
    84  		ch <- l.wait()
    85  	}()
    86  	go func() {
    87  		ch <- cmd.Wait()
    88  	}()
    89  	return <-ch
    90  }
    91  
    92  type listener struct {
    93  	Path string
    94  
    95  	tempDir string
    96  	conn    *net.UnixConn
    97  }
    98  
    99  func listen() (listener, error) {
   100  	dir, err := ioutil.TempDir("", "sdnotify")
   101  	if err != nil {
   102  		return listener{}, err
   103  	}
   104  
   105  	path := filepath.Join(dir, "notify.sock")
   106  	conn, err := net.ListenUnixgram(netType, &net.UnixAddr{
   107  		Net:  netType,
   108  		Name: path,
   109  	})
   110  	if err != nil {
   111  		return listener{}, err
   112  	}
   113  	l := listener{
   114  		Path:    path,
   115  		tempDir: dir,
   116  		conn:    conn}
   117  	return l, nil
   118  }
   119  
   120  func (l listener) wait() error {
   121  	buf := make([]byte, len(readyMsg))
   122  	for {
   123  		n, err := l.conn.Read(buf)
   124  		if err != nil {
   125  			return err
   126  		}
   127  		if string(buf[:n]) == readyMsg {
   128  			return nil
   129  		}
   130  	}
   131  }
   132  
   133  func (l listener) close() error {
   134  	net.Conn(l.conn).Close()
   135  	return os.RemoveAll(l.tempDir)
   136  }