From 8035cf9ed59e74e13e1d882058cb7ef791735aa8 Mon Sep 17 00:00:00 2001
From: Etienne Perot <eperot@google.com>
Date: Mon, 27 Feb 2023 15:24:56 -0800
Subject: [PATCH] `runsc metric-server`: Add `--pid-file` flag for PID file.

This causes the metric server to write its own PID to the given file.
The file is deleted when the metric server exits (cleanly).

Doing so is useful for the usual reasons PID files are useful for, i.e.
identification by other systems that the metric server is running, and knowing
where to send signals.

The PID is also available over HTTP at `/runsc-metrics/pid`.

PiperOrigin-RevId: 512747410
---
 runsc/cmd/metric_server.go | 38 ++++++++++++++++++++++++++++++++++----
 1 file changed, 34 insertions(+), 4 deletions(-)

diff --git a/runsc/cmd/metric_server.go b/runsc/cmd/metric_server.go
index 8f6758d99..d3463242e 100644
--- a/runsc/cmd/metric_server.go
+++ b/runsc/cmd/metric_server.go
@@ -233,6 +233,8 @@ func queryMetrics(ctx context.Context, sand *sandbox.Sandbox, verifier *promethe
 // MetricServer implements subcommands.Command for the "metric-server" command.
 type MetricServer struct {
 	rootDir                string
+	pid                    int
+	pidFile                string
 	allowUnknownRoot       bool
 	exposeProfileEndpoints bool
 	address                string
@@ -296,6 +298,7 @@ func (*MetricServer) Usage() string {
 
 // SetFlags implements subcommands.Command.SetFlags.
 func (m *MetricServer) SetFlags(f *flag.FlagSet) {
+	f.StringVar(&m.pidFile, "pid-file", "", "If set, write the metric server's own PID to this file after binding to the --metric-server address. The parent directory of this file must already exist.")
 	f.BoolVar(&m.exposeProfileEndpoints, "allow-profiling", false, "If true, expose /runsc-metrics/profile-cpu and /runsc-metrics/profile-heap to get profiling data about the metric server")
 }
 
@@ -724,7 +727,7 @@ func (m *MetricServer) serveHealthCheck(w http.ResponseWriter, req *http.Request
 	m.mu.Lock()
 	defer m.mu.Unlock()
 	if m.shuttingDown {
-		return httpResult{http.StatusServiceUnavailable, errors.New("server is shutting down already")}
+		return httpResult{http.StatusServiceUnavailable, errors.New("server is shutting down")}
 	}
 	if err := req.ParseForm(); err != nil {
 		return httpResult{http.StatusBadRequest, err}
@@ -738,6 +741,17 @@ func (m *MetricServer) serveHealthCheck(w http.ResponseWriter, req *http.Request
 	return httpOK
 }
 
+// servePID serves the PID of the metric server process.
+func (m *MetricServer) servePID(w http.ResponseWriter, req *http.Request) httpResult {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	if m.shuttingDown {
+		return httpResult{http.StatusServiceUnavailable, errors.New("server is shutting down")}
+	}
+	io.WriteString(w, strconv.Itoa(m.pid))
+	return httpOK
+}
+
 // profileCPU returns a CPU profile over HTTP.
 func (m *MetricServer) profileCPU(w http.ResponseWriter, req *http.Request) httpResult {
 	// Time to finish up profiling and flush out the results to the client.
@@ -806,6 +820,11 @@ func (m *MetricServer) shutdownLocked(ctx context.Context) {
 			m.udsPath = ""
 		}
 	}
+	if m.pidFile != "" {
+		if err := os.Remove(m.pidFile); err != nil {
+			log.Warningf("Cannot remove PID file at %s: %v", m.pidFile, err)
+		}
+	}
 	m.srv.Shutdown(ctx)
 }
 
@@ -907,6 +926,9 @@ func (m *MetricServer) Execute(ctx context.Context, f *flag.FlagSet, args ...any
 	m.address = conf.MetricServer
 	m.sandboxes = make(map[container.FullID]*servedSandbox)
 	m.lastStateFileStat = make(map[container.FullID]os.FileInfo)
+	m.pid = os.Getpid()
+	m.shutdownCh = make(chan os.Signal, 1)
+	signal.Notify(m.shutdownCh, syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT)
 
 	var listener net.Listener
 	var listenErr error
@@ -951,6 +973,7 @@ func (m *MetricServer) Execute(ctx context.Context, f *flag.FlagSet, args ...any
 
 	mux := http.NewServeMux()
 	mux.HandleFunc("/runsc-metrics/healthcheck", logRequest(m.serveHealthCheck))
+	mux.HandleFunc("/runsc-metrics/pid", logRequest(m.servePID))
 	if m.exposeProfileEndpoints {
 		log.Warningf("Profiling HTTP endpoints are exposed; this should only be used for development!")
 		mux.HandleFunc("/runsc-metrics/profile-cpu", logRequest(m.profileCPU))
@@ -961,11 +984,18 @@ func (m *MetricServer) Execute(ctx context.Context, f *flag.FlagSet, args ...any
 	m.srv.Handler = mux
 	m.srv.ReadTimeout = httpTimeout
 	m.srv.WriteTimeout = httpTimeout
-	m.shutdownCh = make(chan os.Signal, 1)
+	go m.verifyLoop(ctx)
+	if m.pidFile != "" {
+		if err := ioutil.WriteFile(m.pidFile, []byte(fmt.Sprintf("%d", m.pid)), 0644); err != nil {
+			return util.Errorf("Cannot write PID to file %q: %v", m.pidFile, err)
+		}
+		defer os.Remove(m.pidFile)
+		log.Infof("Wrote PID %d to file %v.", m.pid, m.pidFile)
+	}
+
+	// Initialization complete.
 	log.Infof("Server serving on %s for root directory %s.", conf.MetricServer, conf.RootDir)
 
-	go m.verifyLoop(ctx)
-	signal.Notify(m.shutdownCh, syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT)
 	serveErr := m.srv.Serve(listener)
 	log.Infof("Server has stopped accepting requests.")
 	m.mu.Lock()