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

As well as add a test to prove StopUnit is reentrant

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,9 +44,11 @@ func (c *Conn) jobComplete(signal *dbus.Signal) {
 
 	_ = dbus.Store(signal.Body, &id, &job, &unit, &result)
 	c.jobListener.Lock()
-	out, ok := c.jobListener.jobs[job]
+	outs, ok := c.jobListener.jobs[job]
 	if ok {
-		out <- result
+		for _, out := range outs {
+			out <- result
+		}
 		delete(c.jobListener.jobs, job)
 	}
 	c.jobListener.Unlock()
@@ -65,7 +67,10 @@ func (c *Conn) startJob(ctx context.Context, ch chan<- string, job string, args
 	}
 
 	if ch != nil {
-		c.jobListener.jobs[p] = ch
+		if _, ok := c.jobListener.jobs[p]; !ok {
+			c.jobListener.jobs[p] = make([]chan<- string, 0, 1)
+		}
+		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 i := 0; i < 2; i++ {
+		if err := <-errChan; err != nil {
+			t.Fatal(err)
+		}
+	}
+
+	currentJobSize := len(conn.jobListener.jobs)
+	if jobSize != currentJobSize {
+		t.Fatal("JobListener jobs leaked")
+	}
+}