dbus: allow multiple calls for the same unit to *Unit

As well as add a test to prove StopUnit is reentrant

There are cases that multiple StopUnits may be called
for the same job. For instance, in kubernetes, the Kubelet does so in some cases,
when it races against itself to delete a pod. Instead of fixing in the caller, this
can be worked around in this library. Instead of losing track of an old job if a new
StopUnit is called for the same job, we can just track multiple jobs, and signal to
multiple channels when the job finishes.

see kubernetes/kubernetes#135826 and kubernetes/kubernetes#135825
for some more context

Signed-off-by: Peter Hunt <[email protected]>

Author: haircommander

dbus/dbus.go

@@ -95,7 +95,7 @@ type Conn struct {
 	sigobj  dbus.BusObject
 
 	jobListener struct {
-		jobs map[dbus.ObjectPath]chan<- string
+		jobs map[dbus.ObjectPath][]chan<- string
 		sync.Mutex
 	}
 	subStateSubscriber struct {
@@ -207,7 +207,7 @@ func NewConnection(dialBus func() (*dbus.Conn, error)) (*Conn, error) {
 	}
 
 	c.subStateSubscriber.ignore = make(map[dbus.ObjectPath]int64)
-	c.jobListener.jobs = make(map[dbus.ObjectPath]chan<- string)
+	c.jobListener.jobs = make(map[dbus.ObjectPath][]chan<- string)
 
 	// Setup the listeners on jobs so that we can get completions
 	c.sigconn.BusObject().Call("org.freedesktop.DBus.AddMatch", 0,

dbus/methods.go

@@ -44,11 +44,10 @@ func (c *Conn) jobComplete(signal *dbus.Signal) {
 
 	_ = dbus.Store(signal.Body, &id, &job, &unit, &result)
 	c.jobListener.Lock()
-	out, ok := c.jobListener.jobs[job]
-	if ok {
+	for _, out := range c.jobListener.jobs[job] {
 		out <- result
-		delete(c.jobListener.jobs, job)
 	}
+	delete(c.jobListener.jobs, job)
 	c.jobListener.Unlock()
 }
 
@@ -65,7 +64,7 @@ func (c *Conn) startJob(ctx context.Context, ch chan<- string, job string, args
 	}
 
 	if ch != nil {
-		c.jobListener.jobs[p] = ch
+		c.jobListener.jobs[p] = append(c.jobListener.jobs[p], ch)
 	}
 
 	// ignore error since 0 is fine if conversion fails

dbus/methods_test.go

@@ -1749,3 +1749,49 @@ func TestAttachProcessesToUnit(t *testing.T) {
 func TestAttachProcessesToUnitWithSubcgroup(t *testing.T) {
 	testAttachProcessesToUnit(t, "/test-subcgroup")
 }
+
+func TestStopUnitReentrant(t *testing.T) {
+	target := "start-stop.service"
+	conn := setupConn(t)
+
+	setupUnit(target, conn, t)
+	linkUnit(target, conn, t)
+
+	jobSize := len(conn.jobListener.jobs)
+
+	reschan := make(chan string)
+	// Buffered channels are important for multiple calls to the same job,
+	// so the order of when we pull them out doesn't need to matter.
+	reschan2 := make(chan string, 1)
+	reschan3 := make(chan string, 1)
+	errChan := make(chan error, 2)
+	_, err := conn.StartUnit(target, "replace", reschan)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	<-reschan
+
+	go func() {
+		_, err := conn.StopUnit(target, "replace", reschan2)
+		errChan <- err
+	}()
+	go func() {
+		_, err := conn.StopUnit(target, "replace", reschan3)
+		errChan <- err
+	}()
+
+	<-reschan2
+	<-reschan3
+
+	for range 2 {
+		if err := <-errChan; err != nil {
+			t.Fatal(err)
+		}
+	}
+
+	currentJobSize := len(conn.jobListener.jobs)
+	if jobSize != currentJobSize {
+		t.Fatal("JobListener jobs leaked")
+	}
+}