daemon: introduce support for cleaning environment

lucab · lucab · commit 0c088eaedf43 · 2016-11-11T17:11:10.000Z
This introduces an `unsetEnvironment` flag for watchdog and notify
functions to clear the relevant environment variables, aligning those
to their native counterparts.
diff --git a/daemon/sdnotify.go b/daemon/sdnotify.go
@@ -22,11 +22,14 @@ import (
 )
 
 // SdNotify sends a message to the init daemon. It is common to ignore the error.
+// If `unsetEnvironment` is true, the environment variable `NOTIFY_SOCKET`
+// will be unconditionally unset.
+//
 // It returns one of the following:
 // (false, nil) - notification not supported (i.e. NOTIFY_SOCKET is unset)
 // (false, err) - notification supported, but failure happened (e.g. error connecting to NOTIFY_SOCKET or while sending data)
 // (true, nil) - notification supported, data has been sent
-func SdNotify(state string) (sent bool, err error) {
+func SdNotify(unsetEnvironment bool, state string) (sent bool, err error) {
 	socketAddr := &net.UnixAddr{
 		Name: os.Getenv("NOTIFY_SOCKET"),
 		Net:  "unixgram",
@@ -37,6 +40,13 @@ func SdNotify(state string) (sent bool, err error) {
 		return false, nil
 	}
 
+	if unsetEnvironment {
+		err = os.Unsetenv("NOTIFY_SOCKET")
+	}
+	if err != nil {
+		return false, err
+	}
+
 	conn, err := net.DialUnix(socketAddr.Net, nil, socketAddr)
 	// Error connecting to NOTIFY_SOCKET
 	if err != nil {
diff --git a/daemon/sdnotify_test.go b/daemon/sdnotify_test.go
@@ -15,6 +15,7 @@
 package daemon
 
 import (
+	"fmt"
 	"io/ioutil"
 	"net"
 	"os"
@@ -23,9 +24,6 @@ import (
 
 // TestSdNotify
 func TestSdNotify(t *testing.T) {
-	notificationSupportedDataSent := "Notification supported, data sent"
-	notificationSupportedFailure := "Notification supported, but failure happened"
-	notificationNotSupported := "Notification not supported"
 
 	testDir, e := ioutil.TempDir("/tmp/", "test-")
 	if e != nil {
@@ -43,33 +41,39 @@ func TestSdNotify(t *testing.T) {
 		panic(e)
 	}
 
-	// (true, nil) - notification supported, data has been sent
-	e = os.Setenv("NOTIFY_SOCKET", notifySocket)
-	if e != nil {
-		panic(e)
-	}
-	sent, err := SdNotify(notificationSupportedDataSent)
-	if !sent || err != nil {
-		t.Errorf("TEST: %s FAILED", notificationSupportedDataSent)
-	}
+	tests := []struct {
+		unsetEnv  bool
+		envSocket string
 
-	// (false, err) - notification supported, but failure happened
-	e = os.Setenv("NOTIFY_SOCKET", testDir+"/not-exist.sock")
-	if e != nil {
-		panic(e)
-	}
-	sent, err = SdNotify(notificationSupportedFailure)
-	if sent && err == nil {
-		t.Errorf("TEST: %s FAILED", notificationSupportedFailure)
+		wsent bool
+		werr  bool
+	}{
+		// (true, nil) - notification supported, data has been sent
+		{false, notifySocket, true, false},
+		// (false, err) - notification supported, but failure happened
+		{true, testDir + "/missing.sock", false, true},
+		// (false, nil) - notification not supported
+		{true, "", false, false},
 	}
 
-	// (false, nil) - notification not supported
-	e = os.Unsetenv("NOTIFY_SOCKET")
-	if e != nil {
-		panic(e)
-	}
-	sent, err = SdNotify(notificationNotSupported)
-	if sent || err != nil {
-		t.Errorf("TEST: %s FAILED", notificationNotSupported)
+	for i, tt := range tests {
+		must(os.Unsetenv("NOTIFY_SOCKET"))
+		if tt.envSocket != "" {
+			must(os.Setenv("NOTIFY_SOCKET", tt.envSocket))
+		}
+		sent, err := SdNotify(tt.unsetEnv, fmt.Sprintf("TestSdNotify test message #%d", i))
+
+		if sent != tt.wsent {
+			t.Errorf("#%d: expected send result %t, got %t", i, tt.wsent, sent)
+		}
+		if tt.werr && err == nil {
+			t.Errorf("#%d: want non-nil err, got nil", i)
+		} else if !tt.werr && err != nil {
+			t.Errorf("#%d: want nil err, got %v", i, err)
+		}
+		if tt.unsetEnv && tt.envSocket != "" && os.Getenv("NOTIFY_SOCKET") != "" {
+			t.Errorf("#%d: environment variable not cleaned up", i)
+		}
+
 	}
 }
diff --git a/daemon/watchdog.go b/daemon/watchdog.go
@@ -23,14 +23,28 @@ import (
 
 // SdWatchdogEnabled return watchdog information for a service.
 // Process should send daemon.SdNotify("WATCHDOG=1") every time / 2.
+// If `unsetEnvironment` is true, the environment variables `WATCHDOG_USEC`
+// and `WATCHDOG_PID` will be unconditionally unset.
 //
 // It returns one of the following:
 // (0, nil) - watchdog isn't enabled or we aren't the watched PID.
 // (0, err) - an error happened (e.g. error converting time).
 // (time, nil) - watchdog is enabled and we can send ping.
 //   time is delay before inactive service will be killed.
-func SdWatchdogEnabled() (time.Duration, error) {
+func SdWatchdogEnabled(unsetEnvironment bool) (time.Duration, error) {
 	wusec := os.Getenv("WATCHDOG_USEC")
+	wpid := os.Getenv("WATCHDOG_PID")
+	if unsetEnvironment {
+		wusecErr := os.Unsetenv("WATCHDOG_USEC")
+		wpidErr := os.Unsetenv("WATCHDOG_PID")
+		if wusecErr != nil {
+			return 0, wusecErr
+		}
+		if wpidErr != nil {
+			return 0, wpidErr
+		}
+	}
+
 	if wusec == "" {
 		return 0, nil
 	}
@@ -43,7 +57,6 @@ func SdWatchdogEnabled() (time.Duration, error) {
 	}
 	interval := time.Duration(s) * time.Microsecond
 
-	wpid := os.Getenv("WATCHDOG_PID")
 	if wpid == "" {
 		return interval, nil
 	}
diff --git a/daemon/watchdog_test.go b/daemon/watchdog_test.go
@@ -30,29 +30,30 @@ func must(err error) {
 func TestSdWatchdogEnabled(t *testing.T) {
 	mypid := strconv.Itoa(os.Getpid())
 	tests := []struct {
-		usec string // empty => unset
-		pid  string // empty => unset
+		usec     string // empty => unset
+		pid      string // empty => unset
+		unsetEnv bool   // arbitrarily set across testcases
 
 		werr   bool
 		wdelay time.Duration
 	}{
 		// Success cases
-		{"100", mypid, false, 100 * time.Microsecond},
-		{"50", mypid, false, 50 * time.Microsecond},
-		{"1", mypid, false, 1 * time.Microsecond},
-		{"1", "", false, 1 * time.Microsecond},
+		{"100", mypid, true, false, 100 * time.Microsecond},
+		{"50", mypid, true, false, 50 * time.Microsecond},
+		{"1", mypid, false, false, 1 * time.Microsecond},
+		{"1", "", true, false, 1 * time.Microsecond},
 
 		// No-op cases
-		{"", mypid, false, 0}, // WATCHDOG_USEC not set
-		{"1", "0", false, 0},  // WATCHDOG_PID doesn't match
-		{"", "", false, 0},    // Both not set
+		{"", mypid, true, false, 0}, // WATCHDOG_USEC not set
+		{"1", "0", false, false, 0}, // WATCHDOG_PID doesn't match
+		{"", "", true, false, 0},    // Both not set
 
 		// Failure cases
-		{"-1", mypid, true, 0},                // Negative USEC
-		{"string", "1", true, 0},              // Non-integer USEC value
-		{"1", "string", true, 0},              // Non-integer PID value
-		{"stringa", "stringb", true, 0},       // E v e r y t h i n g
-		{"-10239", "-eleventythree", true, 0}, // i s   w r o n g
+		{"-1", mypid, true, true, 0},                // Negative USEC
+		{"string", "1", false, true, 0},             // Non-integer USEC value
+		{"1", "string", true, true, 0},              // Non-integer PID value
+		{"stringa", "stringb", false, true, 0},      // E v e r y t h i n g
+		{"-10239", "-eleventythree", true, true, 0}, // i s   w r o n g
 	}
 
 	for i, tt := range tests {
@@ -67,7 +68,7 @@ func TestSdWatchdogEnabled(t *testing.T) {
 			must(os.Unsetenv("WATCHDOG_PID"))
 		}
 
-		delay, err := SdWatchdogEnabled()
+		delay, err := SdWatchdogEnabled(tt.unsetEnv)
 
 		if tt.werr && err == nil {
 			t.Errorf("#%d: want non-nil err, got nil", i)
@@ -77,5 +78,8 @@ func TestSdWatchdogEnabled(t *testing.T) {
 		if tt.wdelay != delay {
 			t.Errorf("#%d: want delay=%d, got %d", i, tt.wdelay, delay)
 		}
+		if tt.unsetEnv && (os.Getenv("WATCHDOG_PID") != "" || os.Getenv("WATCHDOG_USEC") != "") {
+			t.Errorf("#%d: environment variables not cleaned up", i)
+		}
 	}
 }
