From 88f95ec409eeb6a2d1330a836b0d1b960a3b5c31 Mon Sep 17 00:00:00 2001
From: Brian Downs <brian.downs@gmail.com>
Date: Tue, 15 Jun 2021 07:20:26 -0400
Subject: [PATCH] Send systemd notifications for both server and agent (#3430)

* update agent to sent systemd notify after everything starts
---
 pkg/agent/run.go         | 9 ++++++++-
 pkg/cli/server/server.go | 6 ++----
 2 files changed, 10 insertions(+), 5 deletions(-)

diff --git a/pkg/agent/run.go b/pkg/agent/run.go
index 3cbf3c6240..1c1d5f8f04 100644
--- a/pkg/agent/run.go
+++ b/pkg/agent/run.go
@@ -114,6 +114,10 @@ func run(ctx context.Context, cfg cmds.Agent, proxy proxy.Proxy) error {
 			return err
 		}
 	}
+
+	notifySocket := os.Getenv("NOTIFY_SOCKET")
+	os.Unsetenv("NOTIFY_SOCKET")
+
 	if err := setupTunnelAndRunAgent(ctx, nodeConfig, cfg, proxy); err != nil {
 		return err
 	}
@@ -141,6 +145,9 @@ func run(ctx context.Context, cfg cmds.Agent, proxy proxy.Proxy) error {
 		}
 	}
 
+	os.Setenv("NOTIFY_SOCKET", notifySocket)
+	systemd.SdNotify(true, "READY=1\n")
+
 	<-ctx.Done()
 	return ctx.Err()
 }
@@ -232,7 +239,7 @@ func Run(ctx context.Context, cfg cmds.Agent) error {
 		cfg.Token = newToken.String()
 		break
 	}
-	systemd.SdNotify(true, "READY=1\n")
+
 	return run(ctx, cfg, proxy)
 }
 
diff --git a/pkg/cli/server/server.go b/pkg/cli/server/server.go
index bba86a1b19..b865676af2 100644
--- a/pkg/cli/server/server.go
+++ b/pkg/cli/server/server.go
@@ -362,8 +362,6 @@ func run(app *cli.Context, cfg *cmds.Server, leaderControllers server.CustomCont
 	}
 
 	logrus.Info("Starting " + version.Program + " " + app.App.Version)
-	notifySocket := os.Getenv("NOTIFY_SOCKET")
-	os.Unsetenv("NOTIFY_SOCKET")
 
 	ctx := signals.SetupSignalHandler(context.Background())
 
@@ -379,9 +377,9 @@ func run(app *cli.Context, cfg *cmds.Server, leaderControllers server.CustomCont
 			<-serverConfig.ControlConfig.Runtime.ETCDReady
 			logrus.Info("ETCD server is now running")
 		}
+
 		logrus.Info(version.Program + " is up and running")
-		if notifySocket != "" {
-			os.Setenv("NOTIFY_SOCKET", notifySocket)
+		if cfg.DisableAgent && os.Getenv("NOTIFY_SOCKET") != "" {
 			systemd.SdNotify(true, "READY=1\n")
 		}
 	}()