Repository: pion/webrtc
Pull request #1660 (2021-02-03 09:37:42)

Add SetReadDeadline to RTPSender
#### Description

This PR aims to add SetReadDeadline to RTPSender. Please feel free to let me know if anything is missed. (And yes, I just realized that there is also a `ReadRTCP` method in `RTPSender`.

codecov[bot]: # [Codecov](https://codecov.io/gh/pion/webrtc/pull/1660?src=pr&el=h1) Report
> Merging [#1660](https://codecov.io/gh/pion/webrtc/pull/1660?src=pr&el=desc) (4899500) into [master](https://codecov.io/gh/pion/webrtc/commit/285fb782ba71357baec71a927ea855a86bdd4128?el=desc) (285fb78) will **decrease** coverage by `5.18%`.
> The diff coverage is `n/a`.

[![Impacted file tree graph](https://codecov.io/gh/pion/webrtc/pull/1660/graphs/tree.svg?width=650&height=150&src=pr&token=XOOuQo07I2)](https://codecov.io/gh/pion/webrtc/pull/1660?src=pr&el=tree)

```diff
@@            Coverage Diff             @@
##           master    #1660      +/-   ##
==========================================
- Coverage   75.11%   69.93%   -5.19%     
==========================================
  Files          79       57      -22     
  Lines        5694     2145    -3549     
==========================================
- Hits         4277     1500    -2777     
+ Misses       1037      542     -495     
+ Partials      380      103     -277     
```

| Flag | Coverage Δ | |
|---|---|---|
| go | `?` | |
| wasm | `69.93% <ø> (ø)` | |

Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags#carryforward-flags-in-the-pull-request-comment) to find out more.

| [Impacted Files](https://codecov.io/gh/pion/webrtc/pull/1660?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [icetcp.go](https://codecov.io/gh/pion/webrtc/pull/1660/diff?src=pr&el=tree#diff-aWNldGNwLmdv) | `0.00% <0.00%> (-100.00%)` | :arrow_down: |
| [rtpcodec.go](https://codecov.io/gh/pion/webrtc/pull/1660/diff?src=pr&el=tree#diff-cnRwY29kZWMuZ28=) | `0.00% <0.00%> (-89.48%)` | :arrow_down: |
| [icecandidatepair.go](https://codecov.io/gh/pion/webrtc/pull/1660/diff?src=pr&el=tree#diff-aWNlY2FuZGlkYXRlcGFpci5nbw==) | `0.00% <0.00%> (-85.72%)` | :arrow_down: |
| [track\_local.go](https://codecov.io/gh/pion/webrtc/pull/1660/diff?src=pr&el=tree#diff-dHJhY2tfbG9jYWwuZ28=) | `0.00% <0.00%> (-80.00%)` | :arrow_down: |
| [stats.go](https://codecov.io/gh/pion/webrtc/pull/1660/diff?src=pr&el=tree#diff-c3RhdHMuZ28=) | `0.00% <0.00%> (-67.86%)` | :arrow_down: |
| [internal/mux/muxfunc.go](https://codecov.io/gh/pion/webrtc/pull/1660/diff?src=pr&el=tree#diff-aW50ZXJuYWwvbXV4L211eGZ1bmMuZ28=) | `0.00% <0.00%> (-47.83%)` | :arrow_down: |
| [gathering\_complete\_promise.go](https://codecov.io/gh/pion/webrtc/pull/1660/diff?src=pr&el=tree#diff-Z2F0aGVyaW5nX2NvbXBsZXRlX3Byb21pc2UuZ28=) | `60.00% <0.00%> (-40.00%)` | :arrow_down: |
| [internal/mux/endpoint.go](https://codecov.io/gh/pion/webrtc/pull/1660/diff?src=pr&el=tree#diff-aW50ZXJuYWwvbXV4L2VuZHBvaW50Lmdv) | `22.22% <0.00%> (-38.89%)` | :arrow_down: |
| [sdpsemantics.go](https://codecov.io/gh/pion/webrtc/pull/1660/diff?src=pr&el=tree#diff-c2Rwc2VtYW50aWNzLmdv) | `0.00% <0.00%> (-37.50%)` | :arrow_down: |
| [dtlsrole.go](https://codecov.io/gh/pion/webrtc/pull/1660/diff?src=pr&el=tree#diff-ZHRsc3JvbGUuZ28=) | `68.96% <0.00%> (-31.04%)` | :arrow_down: |
| ... and [29 more](https://codecov.io/gh/pion/webrtc/pull/1660/diff?src=pr&el=tree-more) | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/pion/webrtc/pull/1660?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/pion/webrtc/pull/1660?src=pr&el=footer). Last update [285fb78...95553c3](https://codecov.io/gh/pion/webrtc/pull/1660?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).

Sean-Der: This LGTM @mission-liao!

Would you mind porting over existing tests to your new helper? We have a few tests that wait for PeerConnectionStateConnected and PeerConnectionStateDisconnected would be nice to delete duplicated code :) 

If you don't have the time totally understand! 
mission-liao: @Sean-Der OK, I'm just unsure if there is a guideline to minimize the affected area for each PR. Would love to do that.


Patch:
diff --git a/rtpsender.go b/rtpsender.go
index c29cbeed..f3d34c78 100644
--- a/rtpsender.go
+++ b/rtpsender.go
@@ -5,6 +5,7 @@ package webrtc
 import (
 	"io"
 	"sync"
+	"time"
 
 	"github.com/pion/interceptor"
 	"github.com/pion/randutil"
@@ -233,6 +234,12 @@ func (r *RTPSender) ReadRTCP() ([]rtcp.Packet, interceptor.Attributes, error) {
 	return pkts, attributes, nil
 }
 
+// SetReadDeadline sets the deadline for the Read operation.
+// Setting to zero means no deadline.
+func (r *RTPSender) SetReadDeadline(t time.Time) error {
+	return r.srtpStream.SetReadDeadline(t)
+}
+
 // hasSent tells if data has been ever sent for this instance
 func (r *RTPSender) hasSent() bool {
 	select {
diff --git a/srtp_writer_future.go b/srtp_writer_future.go
index fd3d2210..4d8bafe1 100644
--- a/srtp_writer_future.go
+++ b/srtp_writer_future.go
@@ -5,6 +5,7 @@ package webrtc
 import (
 	"io"
 	"sync/atomic"
+	"time"
 
 	"github.com/pion/rtp"
 	"github.com/pion/srtp/v2"
@@ -80,6 +81,18 @@ func (s *srtpWriterFuture) Read(b []byte) (n int, err error) {
 	return s.Read(b)
 }
 
+func (s *srtpWriterFuture) SetReadDeadline(t time.Time) error {
+	if value := s.rtcpReadStream.Load(); value != nil {
+		return value.(*srtp.ReadStreamSRTCP).SetReadDeadline(t)
+	}
+
+	if err := s.init(false); err != nil || s.rtcpReadStream.Load() == nil {
+		return err
+	}
+
+	return s.SetReadDeadline(t)
+}
+
 func (s *srtpWriterFuture) WriteRTP(header *rtp.Header, payload []byte) (int, error) {
 	if value := s.rtpWriteStream.Load(); value != nil {
 		return value.(*srtp.WriteStreamSRTP).WriteRTP(header, payload)


Test patch:
diff --git a/dtlstransport_test.go b/dtlstransport_test.go
index a4efb493..2bb32e7d 100644
--- a/dtlstransport_test.go
+++ b/dtlstransport_test.go
@@ -3,7 +3,6 @@
 package webrtc
 
 import (
-	"context"
 	"regexp"
 	"testing"
 	"time"
@@ -14,6 +13,9 @@ import (
 
 // An invalid fingerprint MUST cause PeerConnectionState to go to PeerConnectionStateFailed
 func TestInvalidFingerprintCausesFailed(t *testing.T) {
+	lim := test.TimeOut(time.Second * 40)
+	defer lim.Stop()
+
 	report := test.CheckRoutines(t)
 	defer report()
 
@@ -36,12 +38,7 @@ func TestInvalidFingerprintCausesFailed(t *testing.T) {
 		}
 	})
 
-	connectionHasFailed, closeFunc := context.WithCancel(context.Background())
-	pcAnswer.OnConnectionStateChange(func(connectionState PeerConnectionState) {
-		if connectionState == PeerConnectionStateFailed {
-			closeFunc()
-		}
-	})
+	connectionHasFailed := untilConnectionState(PeerConnectionStateFailed, pcAnswer)
 
 	if _, err = pcOffer.CreateDataChannel("unusedDataChannel", nil); err != nil {
 		t.Fatal(err)
@@ -81,11 +78,7 @@ func TestInvalidFingerprintCausesFailed(t *testing.T) {
 		t.Fatal("timed out waiting to receive offer")
 	}
 
-	select {
-	case <-connectionHasFailed.Done():
-	case <-time.After(30 * time.Second):
-		t.Fatal("timed out waiting for connection to fail")
-	}
+	connectionHasFailed.Wait()
 }
 
 func TestPeerConnection_DTLSRoleSettingEngine(t *testing.T) {
@@ -107,18 +100,8 @@ func TestPeerConnection_DTLSRoleSettingEngine(t *testing.T) {
 			t.Fatal(err)
 		}
 
-		connectionComplete := make(chan interface{})
-		answerPC.OnConnectionStateChange(func(connectionState PeerConnectionState) {
-			if connectionState == PeerConnectionStateConnected {
-				select {
-				case <-connectionComplete:
-				default:
-					close(connectionComplete)
-				}
-			}
-		})
-
-		<-connectionComplete
+		connectionComplete := untilConnectionState(PeerConnectionStateConnected, answerPC)
+		connectionComplete.Wait()
 		assert.NoError(t, offerPC.Close())
 		assert.NoError(t, answerPC.Close())
 	}
diff --git a/peerconnection_media_test.go b/peerconnection_media_test.go
index de7c2159..4ebd375d 100644
--- a/peerconnection_media_test.go
+++ b/peerconnection_media_test.go
@@ -1050,17 +1050,7 @@ func TestPeerConnection_Simulcast_Probe(t *testing.T) {
 
 	assert.NoError(t, signalPair(offerer, answerer))
 
-	peerConnectionConnected := sync.WaitGroup{}
-	peerConnectionConnected.Add(2)
-
-	connectionStateHandler := func(connectionState PeerConnectionState) {
-		if connectionState == PeerConnectionStateConnected {
-			peerConnectionConnected.Done()
-		}
-	}
-
-	offerer.OnConnectionStateChange(connectionStateHandler)
-	answerer.OnConnectionStateChange(connectionStateHandler)
+	peerConnectionConnected := untilConnectionState(PeerConnectionStateConnected, offerer, answerer)
 	peerConnectionConnected.Wait()
 
 	<-seenFiveStreams.Done()
diff --git a/rtpreceiver_test.go b/rtpreceiver_test.go
index 3864dc0f..5816bf87 100644
--- a/rtpreceiver_test.go
+++ b/rtpreceiver_test.go
@@ -4,7 +4,6 @@ package webrtc
 
 import (
 	"context"
-	"sync"
 	"testing"
 	"time"
 
@@ -16,7 +15,7 @@ import (
 
 // Assert that SetReadDeadline works as expected
 // This test uses VNet since we must have zero loss
-func Test_RTPSender_SetReadDeadline(t *testing.T) {
+func Test_RTPReceiver_SetReadDeadline(t *testing.T) {
 	lim := test.TimeOut(time.Second * 30)
 	defer lim.Stop()
 
@@ -50,16 +49,7 @@ func Test_RTPSender_SetReadDeadline(t *testing.T) {
 		seenPacketCancel()
 	})
 
-	var peerConnectionsConnected sync.WaitGroup
-	peerConnectionsConnected.Add(2)
-
-	hdlr := func(p PeerConnectionState) {
-		if p == PeerConnectionStateConnected {
-			peerConnectionsConnected.Done()
-		}
-	}
-	sender.OnConnectionStateChange(hdlr)
-	receiver.OnConnectionStateChange(hdlr)
+	peerConnectionsConnected := untilConnectionState(PeerConnectionStateConnected, sender, receiver)
 
 	assert.NoError(t, signalPair(sender, receiver))
 
diff --git a/rtpsender_test.go b/rtpsender_test.go
index f8ac5124..ea1262b1 100644
--- a/rtpsender_test.go
+++ b/rtpsender_test.go
@@ -7,15 +7,32 @@ import (
 	"context"
 	"errors"
 	"io"
+	"sync"
 	"sync/atomic"
 	"testing"
 	"time"
 
+	"github.com/pion/transport/packetio"
 	"github.com/pion/transport/test"
 	"github.com/pion/webrtc/v3/pkg/media"
 	"github.com/stretchr/testify/assert"
 )
 
+func untilConnectionState(state PeerConnectionState, peers ...*PeerConnection) *sync.WaitGroup {
+	var triggered sync.WaitGroup
+	triggered.Add(len(peers))
+
+	hdlr := func(p PeerConnectionState) {
+		if p == state {
+			triggered.Done()
+		}
+	}
+	for _, p := range peers {
+		p.OnConnectionStateChange(hdlr)
+	}
+	return &triggered
+}
+
 func Test_RTPSender_ReplaceTrack(t *testing.T) {
 	lim := test.TimeOut(time.Second * 10)
 	defer lim.Stop()
@@ -158,3 +175,33 @@ func Test_RTPSender_GetParameters(t *testing.T) {
 	assert.NoError(t, offerer.Close())
 	assert.NoError(t, answerer.Close())
 }
+
+func Test_RTPSender_SetReadDeadline(t *testing.T) {
+	lim := test.TimeOut(time.Second * 30)
+	defer lim.Stop()
+
+	report := test.CheckRoutines(t)
+	defer report()
+
+	sender, receiver, wan := createVNetPair(t)
+
+	track, err := NewTrackLocalStaticSample(RTPCodecCapability{MimeType: "video/vp8"}, "video", "pion")
+	assert.NoError(t, err)
+
+	rtpSender, err := sender.AddTrack(track)
+	assert.NoError(t, err)
+
+	peerConnectionsConnected := untilConnectionState(PeerConnectionStateConnected, sender, receiver)
+
+	assert.NoError(t, signalPair(sender, receiver))
+
+	peerConnectionsConnected.Wait()
+
+	assert.NoError(t, rtpSender.SetReadDeadline(time.Now().Add(1*time.Second)))
+	_, _, err = rtpSender.ReadRTCP()
+	assert.Error(t, err, packetio.ErrTimeout)
+
+	assert.NoError(t, wan.Stop())
+	assert.NoError(t, sender.Close())
+	assert.NoError(t, receiver.Close())
+}