fix: quit runwda when wda dies (danielpaulus#389)

With this changes ios runwda will exit with error when WDA app dies (e.g. killed with ios kill, killed manually with AppSwitcher, or when iPhone is turned off).
dtxConnection's BreakdownCallback is implemented with Functionals Options Pattern, so it's fully optional and shouldn't affect any other dtxConnection usages.

Author: michal-przytarski

ios/dtx_codec/connection.go

@@ -1,6 +1,7 @@
 package dtx
 
 import (
+	"errors"
 	"io"
 	"math"
 	"strings"
@@ -15,6 +16,8 @@ import (
 
 type MethodWithResponse func(msg Message) (interface{}, error)
 
+var ErrConnectionClosed = errors.New("Connection closed")
+
 // Connection manages channels, including the GlobalChannel, for a DtxConnection and dispatches received messages
 // to the right channel.
 type Connection struct {
@@ -25,6 +28,10 @@ type Connection struct {
 	capabilities           map[string]interface{}
 	mutex                  sync.Mutex
 	requestChannelMessages chan Message
+
+	closed    chan struct{}
+	err       error
+	closeOnce sync.Once
 }
 
 // Dispatcher is a simple interface containing a Dispatch func to receive dtx.Messages
@@ -41,11 +48,24 @@ type GlobalDispatcher struct {
 
 const requestChannel = "_requestChannelWithCode:identifier:"
 
+// Closed is closed when the underlying DTX connection was closed for any reason (either initiated by calling Close() or due to an error)
+func (dtxConn *Connection) Closed() <-chan struct{} {
+	return dtxConn.closed
+}
+
+// Err is non-nil when the connection was closed (when Close was called this will be ErrConnectionClosed)
+func (dtxConn *Connection) Err() error {
+	return dtxConn.err
+}
+
 // Close closes the underlying deviceConnection
 func (dtxConn *Connection) Close() error {
 	if dtxConn.deviceConnection != nil {
-		return dtxConn.deviceConnection.Close()
+		err := dtxConn.deviceConnection.Close()
+		dtxConn.close(err)
+		return err
 	}
+	dtxConn.close(ErrConnectionClosed)
 	return nil
 }
 
@@ -121,6 +141,7 @@ func newDtxConnection(conn ios.DeviceConnectionInterface) (*Connection, error) {
 
 	// The global channel has channelCode 0, so we need to start with channelCodeCounter==1
 	dtxConnection := &Connection{deviceConnection: conn, channelCodeCounter: 1, requestChannelMessages: requestChannelMessages}
+	dtxConnection.closed = make(chan struct{})
 
 	// The global channel is automatically present and used for requesting other channels and some other methods like notifyPublishedCapabilities
 	globalChannel := Channel{
@@ -149,6 +170,7 @@ func reader(dtxConn *Connection) {
 		reader := dtxConn.deviceConnection.Reader()
 		msg, err := ReadMessage(reader)
 		if err != nil {
+			defer dtxConn.close(err)
 			errText := err.Error()
 			if err == io.EOF || strings.Contains(errText, "use of closed network") {
 				log.Debug("DTX Connection with EOF")
@@ -228,3 +250,10 @@ func (dtxConn *Connection) RequestChannelIdentifier(identifier string, messageDi
 
 	return channel
 }
+
+func (dtxConn *Connection) close(err error) {
+	dtxConn.closeOnce.Do(func() {
+		dtxConn.err = err
+		close(dtxConn.closed)
+	})
+}

ios/testmanagerd/xcuitestrunner.go

@@ -393,16 +393,29 @@ func runXUITestWithBundleIdsXcode15Ctx(
 	}
 
 	select {
+	case <-conn1.Closed():
+		log.Debug("conn1 closed")
+		if conn1.Err() != dtx.ErrConnectionClosed {
+			log.WithError(conn1.Err()).Error("conn1 closed unexpectedly")
+		}
+		break
+	case <-conn2.Closed():
+		log.Debug("conn2 closed")
+		if conn2.Err() != dtx.ErrConnectionClosed {
+			log.WithError(conn2.Err()).Error("conn2 closed unexpectedly")
+		}
+		break
 	case <-testListener.Done():
 		break
 	case <-ctx.Done():
-		log.Infof("Killing test runner with pid %d ...", testRunnerLaunch.Pid)
-		err = killTestRunner(appserviceConn, testRunnerLaunch.Pid)
-		if err != nil {
-			log.Infof("Nothing to kill, process with pid %d is already dead", testRunnerLaunch.Pid)
-		} else {
-			log.Info("Test runner killed with success")
-		}
+		break
+	}
+	log.Infof("Killing test runner with pid %d ...", testRunnerLaunch.Pid)
+	err = killTestRunner(appserviceConn, testRunnerLaunch.Pid)
+	if err != nil {
+		log.Infof("Nothing to kill, process with pid %d is already dead", testRunnerLaunch.Pid)
+	} else {
+		log.Info("Test runner killed with success")
 	}
 
 	log.Debugf("Done running test")

ios/testmanagerd/xcuitestrunner_11.go

@@ -78,16 +78,29 @@ func runXCUIWithBundleIdsXcode11Ctx(
 	}
 
 	select {
+	case <-conn.Closed():
+		log.Debug("conn closed")
+		if conn.Err() != dtx.ErrConnectionClosed {
+			log.WithError(conn.Err()).Error("conn closed unexpectedly")
+		}
+		break
+	case <-conn2.Closed():
+		log.Debug("conn2 closed")
+		if conn2.Err() != dtx.ErrConnectionClosed {
+			log.WithError(conn2.Err()).Error("conn2 closed unexpectedly")
+		}
+		break
 	case <-testListener.Done():
 		break
 	case <-ctx.Done():
-		log.Infof("Killing test runner with pid %d ...", pid)
-		err = pControl.KillProcess(pid)
-		if err != nil {
-			log.Infof("Nothing to kill, process with pid %d is already dead", pid)
-		} else {
-			log.Info("Test runner killed with success")
-		}
+		break
+	}
+	log.Infof("Killing test runner with pid %d ...", pid)
+	err = pControl.KillProcess(pid)
+	if err != nil {
+		log.Infof("Nothing to kill, process with pid %d is already dead", pid)
+	} else {
+		log.Info("Test runner killed with success")
 	}
 
 	log.Debugf("Done running test")

ios/testmanagerd/xcuitestrunner_12.go

@@ -77,16 +77,29 @@ func runXUITestWithBundleIdsXcode12Ctx(ctx context.Context, bundleID string, tes
 	}
 
 	select {
+	case <-conn.Closed():
+		log.Debug("conn closed")
+		if conn.Err() != dtx.ErrConnectionClosed {
+			log.WithError(conn.Err()).Error("conn closed unexpectedly")
+		}
+		break
+	case <-conn2.Closed():
+		log.Debug("conn2 closed")
+		if conn2.Err() != dtx.ErrConnectionClosed {
+			log.WithError(conn2.Err()).Error("conn2 closed unexpectedly")
+		}
+		break
 	case <-testListener.Done():
 		break
 	case <-ctx.Done():
-		log.Infof("Killing test runner with pid %d ...", pid)
-		err = pControl.KillProcess(pid)
-		if err != nil {
-			log.Infof("Nothing to kill, process with pid %d is already dead", pid)
-		} else {
-			log.Info("Test runner killed with success")
-		}
+		break
+	}
+	log.Infof("Killing test runner with pid %d ...", pid)
+	err = pControl.KillProcess(pid)
+	if err != nil {
+		log.Infof("Nothing to kill, process with pid %d is already dead", pid)
+	} else {
+		log.Info("Test runner killed with success")
 	}
 
 	log.Debugf("Done running test")

main.go

@@ -1124,28 +1124,30 @@ func runWdaCommand(device ios.DeviceEntry, arguments docopt.Opts) bool {
 		log.WithFields(log.Fields{"bundleid": bundleID, "testbundleid": testbundleID, "xctestconfig": xctestconfig}).Info("Running wda")
 
 		errorChannel := make(chan error)
+		defer close(errorChannel)
 		ctx, stopWda := context.WithCancel(context.Background())
 		go func() {
 			_, err := testmanagerd.RunXCUIWithBundleIdsCtx(ctx, bundleID, testbundleID, xctestconfig, device, wdaargs, wdaenv, nil, nil, testmanagerd.NewTestListener(io.Discard, io.Discard, os.TempDir()))
 			if err != nil {
-				log.WithFields(log.Fields{"error": err}).Fatal("Failed running WDA")
 				errorChannel <- err
 			}
-			close(errorChannel)
+			stopWda()
 		}()
 		c := make(chan os.Signal, 1)
 		signal.Notify(c, syscall.SIGINT, syscall.SIGTERM)
-		signal := <-c
-		log.Infof("os signal:%d received, closing..", signal)
 
-		stopWda()
-
-		err := <-errorChannel
-		if err != nil {
-			log.Errorf("Failed running wda-testrunner: %s", err)
+		select {
+		case err := <-errorChannel:
+			log.WithError(err).Error("Failed running WDA")
+			stopWda()
 			os.Exit(1)
+		case <-ctx.Done():
+			log.Error("WDA process ended unexpectedly")
+			os.Exit(1)
+		case signal := <-c:
+			log.Infof("os signal:%d received, closing...", signal)
+			stopWda()
 		}
-
 		log.Info("Done Closing")
 	}
 	return b