Refactor timeout loop to use a context and reuse timers

This commit is contained in:
Cecylia Bocovich 2022-11-28 17:30:05 -05:00
parent b010de5abb
commit b547d449cb
No known key found for this signature in database
GPG key ID: 009DE379FD9B7B90

View file

@ -1,6 +1,7 @@
package snowflake_proxy package snowflake_proxy
import ( import (
"context"
"fmt" "fmt"
"io" "io"
"log" "log"
@ -35,29 +36,35 @@ type webRTCConn struct {
inactivityTimeout time.Duration inactivityTimeout time.Duration
activity chan struct{} activity chan struct{}
closed chan struct{} cancelTimeoutLoop context.CancelFunc
} }
func newWebRTCConn(pc *webrtc.PeerConnection, dc *webrtc.DataChannel, pr *io.PipeReader, eventLogger event.SnowflakeEventReceiver) *webRTCConn { func newWebRTCConn(pc *webrtc.PeerConnection, dc *webrtc.DataChannel, pr *io.PipeReader, eventLogger event.SnowflakeEventReceiver) *webRTCConn {
conn := &webRTCConn{pc: pc, dc: dc, pr: pr, eventLogger: eventLogger} conn := &webRTCConn{pc: pc, dc: dc, pr: pr, eventLogger: eventLogger}
conn.bytesLogger = newBytesSyncLogger() conn.bytesLogger = newBytesSyncLogger()
conn.activity = make(chan struct{}, 100) conn.activity = make(chan struct{}, 100)
conn.closed = make(chan struct{})
conn.inactivityTimeout = 30 * time.Second conn.inactivityTimeout = 30 * time.Second
go conn.timeoutLoop() ctx, cancel := context.WithCancel(context.Background())
conn.cancelTimeoutLoop = cancel
go conn.timeoutLoop(ctx)
return conn return conn
} }
func (c *webRTCConn) timeoutLoop() { func (c *webRTCConn) timeoutLoop(ctx context.Context) {
timer := time.NewTimer(c.inactivityTimeout)
for { for {
select { select {
case <-time.After(c.inactivityTimeout): case <-timer.C:
c.Close() c.Close()
log.Println("Closed connection due to inactivity") log.Println("Closed connection due to inactivity")
return return
case <-c.activity: case <-c.activity:
if !timer.Stop() {
<-timer.C
}
timer.Reset(c.inactivityTimeout)
continue continue
case <-c.closed: case <-ctx.Done():
return return
} }
} }
@ -80,7 +87,7 @@ func (c *webRTCConn) Write(b []byte) (int, error) {
func (c *webRTCConn) Close() (err error) { func (c *webRTCConn) Close() (err error) {
c.once.Do(func() { c.once.Do(func() {
close(c.closed) c.cancelTimeoutLoop()
err = c.pc.Close() err = c.pc.Close()
}) })
return return