dji/tello: fix the issue about Halt method

platforms/dji/tello/driver.go

@@ -10,6 +10,7 @@ import (
 	"net"
 	"strconv"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"gobot.io/x/gobot"
@@ -191,7 +192,8 @@ type Driver struct {
 	throttle       int
 	bouncing       bool
 	gobot.Eventer
-	doneCh chan struct{}
+	doneCh            chan struct{}
+	doneChReaderCount int32
 }
 
 // NewDriver creates a driver for the Tello drone. Pass in the UDP port to use for the responses
@@ -280,7 +282,10 @@ func (d *Driver) Start() error {
 	d.cmdConn = cmdConn
 
 	// handle responses
+	d.addDoneChReaderCount(1)
 	go func() {
+		defer d.addDoneChReaderCount(-1)
+
 		d.On(d.Event(ConnectedEvent), func(interface{}) {
 			d.SendDateTime()
 			d.processVideo()
@@ -304,13 +309,22 @@ func (d *Driver) Start() error {
 	d.SendCommand(d.connectionString())
 
 	// send stick commands
+	d.addDoneChReaderCount(1)
 	go func() {
+		defer d.addDoneChReaderCount(-1)
+
+	stickCmdLoop:
 		for {
-			err := d.SendStickCommand()
-			if err != nil {
-				fmt.Println("stick command error:", err)
+			select {
+			case <-d.doneCh:
+				break stickCmdLoop
+			default:
+				err := d.SendStickCommand()
+				if err != nil {
+					fmt.Println("stick command error:", err)
+				}
+				time.Sleep(20 * time.Millisecond)
 			}
-			time.Sleep(20 * time.Millisecond)
 		}
 	}()
 
@@ -320,14 +334,23 @@ func (d *Driver) Start() error {
 // Halt stops the driver.
 func (d *Driver) Halt() (err error) {
 	// send a landing command when we disconnect, and give it 500ms to be received before we shutdown
-	d.Land()
-	d.doneCh <- struct{}{}
+	if d.cmdConn != nil {
+		d.Land()
+	}
 	time.Sleep(500 * time.Millisecond)
 
-	d.cmdConn.Close()
+	if d.cmdConn != nil {
+		d.cmdConn.Close()
+	}
+
 	if d.videoConn != nil {
 		d.videoConn.Close()
 	}
+	readerCount := atomic.LoadInt32(&d.doneChReaderCount)
+	for i := 0; i < int(readerCount); i++ {
+		d.doneCh <- struct{}{}
+	}
+
 	return
 }
 
@@ -946,7 +969,10 @@ func (d *Driver) processVideo() error {
 		return err
 	}
 
+	d.addDoneChReaderCount(1)
 	go func() {
+		defer d.addDoneChReaderCount(-1)
+
 	videoConnLoop:
 		for {
 			select {
@@ -989,6 +1015,10 @@ func (d *Driver) connectionString() string {
 	return res
 }
 
+func (d *Driver) addDoneChReaderCount(delta int32) {
+	atomic.AddInt32(&d.doneChReaderCount, delta)
+}
+
 func (f *FlightData) AirSpeed() float64 {
 	return math.Sqrt(
 		math.Pow(float64(f.NorthSpeed), 2) +

platforms/dji/tello/driver_test.go

@@ -5,6 +5,7 @@ import (
 	"encoding/binary"
 	"fmt"
 	"io"
+	"sync"
 	"testing"
 	"time"
 
@@ -14,6 +15,15 @@ import (
 
 var _ gobot.Driver = (*Driver)(nil)
 
+type WriteCloserDoNothing struct{}
+
+func (w *WriteCloserDoNothing) Write(p []byte) (n int, err error) {
+	return 0, nil
+}
+func (w *WriteCloserDoNothing) Close() error {
+	return nil
+}
+
 func TestTelloDriver(t *testing.T) {
 	d := NewDriver("8888")
 
@@ -134,3 +144,49 @@ func TestHandleResponse(t *testing.T) {
 		})
 	}
 }
+
+func TestHaltShouldTerminateAllTheRelatedGoroutines(t *testing.T) {
+	d := NewDriver("8888")
+	d.cmdConn = &WriteCloserDoNothing{}
+
+	var wg sync.WaitGroup
+	wg.Add(3)
+
+	d.addDoneChReaderCount(1)
+	go func() {
+		<-d.doneCh
+		d.addDoneChReaderCount(-1)
+		wg.Done()
+		fmt.Println("Done routine 1.")
+	}()
+
+	d.addDoneChReaderCount(1)
+	go func() {
+		<-d.doneCh
+		d.addDoneChReaderCount(-1)
+		wg.Done()
+		fmt.Println("Done routine 2.")
+	}()
+
+	d.addDoneChReaderCount(1)
+	go func() {
+		<-d.doneCh
+		d.addDoneChReaderCount(-1)
+		wg.Done()
+		fmt.Println("Done routine 3.")
+	}()
+
+	d.Halt()
+	wg.Wait()
+
+	gobottest.Assert(t, d.doneChReaderCount, int32(0))
+}
+
+func TestHaltNotWaitForeverWhenCalledMultipleTimes(t *testing.T) {
+	d := NewDriver("8888")
+	d.cmdConn = &WriteCloserDoNothing{}
+
+	d.Halt()
+	d.Halt()
+	d.Halt()
+}