From 3b355ef0e5dab14d8ac75248f573e3206fc12c39 Mon Sep 17 00:00:00 2001
From: Ivan Babrou <github@ivan.computer>
Date: Wed, 17 Apr 2024 20:42:04 -0700
Subject: [PATCH] Add sd_notify support
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

When running under systemd with `Type=notify`, the service would only
report as ready after it's fully started and loaded all programs.

Fully running service:

```
● ebpf_exporter.service
     Loaded: loaded (/etc/systemd/system/ebpf_exporter.service; static)
     Active: active (running) since Thu 2024-04-18 03:40:41 UTC; 2min 13s ago
   Main PID: 24569 (ebpf_exporter)
     Status: "ready with 1 programs found in the config in 278ms"
```
---
 cmd/ebpf_exporter/main.go | 49 ++++++++++++++++++++++++++++++++-------
 go.mod                    |  1 +
 go.sum                    |  2 ++
 3 files changed, 43 insertions(+), 9 deletions(-)

diff --git a/cmd/ebpf_exporter/main.go b/cmd/ebpf_exporter/main.go
index 1b23ce17..aa3895c8 100644
--- a/cmd/ebpf_exporter/main.go
+++ b/cmd/ebpf_exporter/main.go
@@ -4,6 +4,7 @@ import (
 	"errors"
 	"fmt"
 	"log"
+	"net"
 	"net/http"
 	_ "net/http/pprof"
 	"os"
@@ -16,6 +17,7 @@ import (
 	"github.com/cloudflare/ebpf_exporter/v2/exporter"
 	"github.com/cloudflare/ebpf_exporter/v2/tracing"
 	"github.com/coreos/go-systemd/activation"
+	"github.com/mdlayher/sdnotify"
 	"github.com/prometheus/client_golang/prometheus"
 	version_collector "github.com/prometheus/client_golang/prometheus/collectors/version"
 	"github.com/prometheus/client_golang/prometheus/promhttp"
@@ -56,8 +58,20 @@ func main() {
 
 	libbpfgo.SetLoggerCbs(libbpfgoCallbacks)
 
+	notifier, _ := sdnotify.New()
+
+	notify := func(format string, v ...interface{}) {
+		if notifier == nil {
+			return
+		}
+
+		_ = notifier.Notify(sdnotify.Statusf(format, v...))
+	}
+
 	started := time.Now()
 
+	notify("parsing config...")
+
 	configs, err := config.ParseConfigs(*configDir, strings.Split(*configNames, ","))
 	if err != nil {
 		log.Fatalf("Error parsing configs: %v", err)
@@ -72,11 +86,15 @@ func main() {
 		log.Fatalf("Error creating tracing processor: %v", err)
 	}
 
+	notify("creating exporter...")
+
 	e, err := exporter.New(configs, tracing.NewProvider(processor), *btfPath)
 	if err != nil {
 		log.Fatalf("Error creating exporter: %s", err)
 	}
 
+	notify("attaching configs...")
+
 	err = e.Attach()
 	if err != nil {
 		log.Fatalf("Error attaching exporter: %s", err)
@@ -87,7 +105,9 @@ func main() {
 		log.Fatalf("Error dropping capabilities: %s", err)
 	}
 
-	log.Printf("Started with %d programs found in the config in %dms", len(configs), time.Since(started).Milliseconds())
+	readyString := fmt.Sprintf("%d programs found in the config in %dms", len(configs), time.Since(started).Milliseconds())
+
+	log.Printf("Started with %s", readyString)
 
 	if *configStrict {
 		missed := e.MissedAttachments()
@@ -130,33 +150,44 @@ func main() {
 		http.HandleFunc("/maps", e.MapsHandler)
 	}
 
-	err = listen(*listenAddress)
+	listener, err := listen(*listenAddress)
+	if err != nil {
+		log.Fatalf("Error listening on %s: %v", *listenAddress, err)
+	}
+
+	notify("ready with %s", readyString)
+
+	if notifier != nil {
+		_ = notifier.Notify(sdnotify.Ready)
+	}
+
+	err = http.Serve(listener, nil)
 	if err != nil {
-		log.Fatalf("Error listening on %s: %s", *listenAddress, err)
+		log.Fatalf("Error serving HTTP: %v", err)
 	}
 }
 
-func listen(addr string) error {
+func listen(addr string) (net.Listener, error) {
 	log.Printf("Listening on %s", addr)
 	if strings.HasPrefix(addr, "fd://") {
 		fd, err := strconv.Atoi(strings.TrimPrefix(addr, "fd://"))
 		if err != nil {
-			return fmt.Errorf("error extracting fd number from %q: %v", addr, err)
+			return nil, fmt.Errorf("error extracting fd number from %q: %v", addr, err)
 		}
 
 		listeners, err := activation.Listeners()
 		if err != nil {
-			return fmt.Errorf("error getting activation listeners: %v", err)
+			return nil, fmt.Errorf("error getting activation listeners: %v", err)
 		}
 
 		if len(listeners) < fd+1 {
-			return errors.New("no listeners passed via activation")
+			return nil, errors.New("no listeners passed via activation")
 		}
 
-		return http.Serve(listeners[fd], nil)
+		return listeners[fd], nil
 	}
 
-	return http.ListenAndServe(addr, nil)
+	return net.Listen("tcp", addr)
 }
 
 func ensureCapabilities(keep string) error {
diff --git a/go.mod b/go.mod
index 4369c6bc..3cd9dba2 100644
--- a/go.mod
+++ b/go.mod
@@ -8,6 +8,7 @@ require (
 	github.com/elastic/go-perf v0.0.0-20191212140718-9c656876f595
 	github.com/iovisor/gobpf v0.2.0
 	github.com/jaypipes/pcidb v1.0.0
+	github.com/mdlayher/sdnotify v1.0.0
 	github.com/prometheus/client_golang v1.19.0
 	github.com/prometheus/common v0.51.1
 	go.opentelemetry.io/contrib/exporters/autoexport v0.49.0
diff --git a/go.sum b/go.sum
index 6ffde22e..400348d0 100644
--- a/go.sum
+++ b/go.sum
@@ -38,6 +38,8 @@ github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE=
 github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk=
 github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=
 github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE=
+github.com/mdlayher/sdnotify v1.0.0 h1:Ma9XeLVN/l0qpyx1tNeMSeTjCPH6NtuD6/N9XdTlQ3c=
+github.com/mdlayher/sdnotify v1.0.0/go.mod h1:HQUmpM4XgYkhDLtd+Uad8ZFK1T9D5+pNxnXQjCeJlGE=
 github.com/mitchellh/go-homedir v1.1.0 h1:lukF9ziXFxDFPkA1vsr5zpc1XuPDn/wFntq5mG+4E0Y=
 github.com/mitchellh/go-homedir v1.1.0/go.mod h1:SfyaCUpYCn1Vlf4IUYiD9fPX4A5wJrkLzIz1N1q0pr0=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=