fix client-publish-udp example

connclient.go

@@ -622,18 +622,18 @@ func (c *ConnClient) Setup(u *base.URL, mode headers.TransportMode, proto base.S
 }
 
 // Play writes a PLAY request and reads a Response.
-// This function can be called only after SetupUDP() or SetupTCP().
+// This can be called only after Setup().
 func (c *ConnClient) Play(u *base.URL) (*base.Response, error) {
 	if c.state != connClientStateInitial {
 		return nil, fmt.Errorf("can't be called when reading or publishing")
 	}
 
 	if c.streamUrl == nil {
-		return nil, fmt.Errorf("can be called only after a successful SetupUDP() or SetupTCP()")
+		return nil, fmt.Errorf("can be called only after a successful Setup()")
 	}
 
 	if *u != *c.streamUrl {
-		return nil, fmt.Errorf("must be called with the same url used for SetupUDP() or SetupTCP()")
+		return nil, fmt.Errorf("must be called with the same url used for Setup())")
 	}
 
 	res, err := c.Do(&base.Request{
@@ -768,7 +768,7 @@ func (c *ConnClient) LoopUDP() error {
 // Announce writes an ANNOUNCE request and reads a Response.
 func (c *ConnClient) Announce(u *base.URL, tracks Tracks) (*base.Response, error) {
 	if c.streamUrl != nil {
-		return nil, fmt.Errorf("announce has already been sent with another url url")
+		return nil, fmt.Errorf("announce has already been sent with another url")
 	}
 
 	res, err := c.Do(&base.Request{
@@ -793,6 +793,7 @@ func (c *ConnClient) Announce(u *base.URL, tracks Tracks) (*base.Response, error
 }
 
 // Record writes a RECORD request and reads a Response.
+// This can be called only after Announce() and Setup().
 func (c *ConnClient) Record(u *base.URL) (*base.Response, error) {
 	if c.state != connClientStateInitial {
 		return nil, fmt.Errorf("can't be called when reading or publishing")

dialer_test.go

@@ -204,14 +204,17 @@ func TestDialPublishUDP(t *testing.T) {
 			defer func() { <-publishDone }()
 
 			var conn *ConnClient
-			defer func() {
-				conn.Close()
-				conn.LoopUDP()
-			}()
 
 			go func() {
 				defer close(publishDone)
 
+				var writerDone chan struct{}
+				defer func() {
+					if writerDone != nil {
+						<-writerDone
+					}
+				}()
+
 				pc, err := net.ListenPacket("udp4", "127.0.0.1:0")
 				require.NoError(t, err)
 				defer pc.Close()
@@ -238,6 +241,11 @@ func TestDialPublishUDP(t *testing.T) {
 					StreamProtocolUDP, Tracks{track})
 				require.NoError(t, err)
 
+				writerDone = make(chan struct{})
+
+				go func() {
+					defer close(writerDone)
+
 					buf := make([]byte, 2048)
 					for {
 						n, _, err := pc.ReadFrom(buf)
@@ -252,6 +260,9 @@ func TestDialPublishUDP(t *testing.T) {
 					}
 				}()
 
+				conn.LoopUDP()
+			}()
+
 			if server == "ffmpeg" {
 				time.Sleep(5 * time.Second)
 			}
@@ -269,6 +280,8 @@ func TestDialPublishUDP(t *testing.T) {
 
 			code := cnt3.wait()
 			require.Equal(t, 0, code)
+
+			conn.Close()
 		})
 	}
 }
@@ -303,9 +316,6 @@ func TestDialPublishTCP(t *testing.T) {
 			defer func() { <-publishDone }()
 
 			var conn *ConnClient
-			defer func() {
-				conn.Close()
-			}()
 
 			go func() {
 				defer close(publishDone)
@@ -367,6 +377,8 @@ func TestDialPublishTCP(t *testing.T) {
 
 			code := cnt3.wait()
 			require.Equal(t, 0, code)
+
+			conn.Close()
 		})
 	}
 }

examples/client-publish-udp.go

@@ -15,6 +15,13 @@ import (
 // the frames with the UDP protocol.
 
 func main() {
+	var writerDone chan struct{}
+	defer func() {
+		if writerDone != nil {
+			<-writerDone
+		}
+	}()
+
 	// open a listener to receive RTP/H264 frames
 	pc, err := net.ListenPacket("udp4", "127.0.0.1:9000")
 	if err != nil {
@@ -48,6 +55,11 @@ func main() {
 	}
 	defer conn.Close()
 
+	writerDone = make(chan struct{})
+
+	go func() {
+		defer close(writerDone)
+
 		buf := make([]byte, 2048)
 		for {
 			// read frames from the source
@@ -62,6 +74,7 @@ func main() {
 				break
 			}
 		}
+	}()
 
 	// wait until the connection is closed
 	err = conn.LoopUDP()