TUN-5488: Close session after it's idle for a period defined by registerUdpSession RPC

datagramsession/session.go

@@ -3,10 +3,15 @@ package datagramsession
 import (
 	"context"
 	"io"
+	"time"
 
 	"github.com/google/uuid"
 )
 
+const (
+	defaultCloseIdleAfter = time.Second * 210
+)
+
 // Each Session is a bidirectional pipe of datagrams between transport and dstConn
 // Currently the only implementation of transport is quic DatagramMuxer
 // Destination can be a connection with origin or with eyeball
@@ -22,7 +27,9 @@ type Session struct {
 	id        uuid.UUID
 	transport transport
 	dstConn   io.ReadWriteCloser
-	doneChan  chan struct{}
+	// activeAtChan is used to communicate the last read/write time
+	activeAtChan chan time.Time
+	doneChan     chan struct{}
 }
 
 func newSession(id uuid.UUID, transport transport, dstConn io.ReadWriteCloser) *Session {
@@ -30,41 +37,81 @@ func newSession(id uuid.UUID, transport transport, dstConn io.ReadWriteCloser) *
 		id:        id,
 		transport: transport,
 		dstConn:   dstConn,
-		doneChan:  make(chan struct{}),
+		// activeAtChan has low capacity. It can be full when there are many concurrent read/write. markActive() will
+		// drop instead of blocking because last active time only needs to be an approximation
+		activeAtChan: make(chan time.Time, 2),
+		doneChan:     make(chan struct{}),
 	}
 }
 
-func (s *Session) Serve(ctx context.Context) error {
+func (s *Session) Serve(ctx context.Context, closeAfterIdle time.Duration) error {
 	serveCtx, cancel := context.WithCancel(ctx)
 	defer cancel()
-	go func() {
-		select {
-		case <-serveCtx.Done():
-		case <-s.doneChan:
-		}
-		s.dstConn.Close()
-	}()
+	go s.waitForCloseCondition(serveCtx, closeAfterIdle)
 	// QUIC implementation copies data to another buffer before returning https://github.com/lucas-clemente/quic-go/blob/v0.24.0/session.go#L1967-L1975
 	// This makes it safe to share readBuffer between iterations
-	readBuffer := make([]byte, 1280)
+	readBuffer := make([]byte, s.transport.MTU())
 	for {
-		// TODO: TUN-5303: origin proxy should determine the buffer size
-		n, err := s.dstConn.Read(readBuffer)
-		if n > 0 {
-			if err := s.transport.SendTo(s.id, readBuffer[:n]); err != nil {
-				return err
-			}
-		}
-		if err != nil {
+		if err := s.dstToTransport(readBuffer); err != nil {
 			return err
 		}
 	}
 }
 
-func (s *Session) writeToDst(payload []byte) (int, error) {
+func (s *Session) waitForCloseCondition(ctx context.Context, closeAfterIdle time.Duration) {
+	if closeAfterIdle == 0 {
+		// provide deafult is caller doesn't specify one
+		closeAfterIdle = defaultCloseIdleAfter
+	}
+	// Closing dstConn cancels read so Serve function can return
+	defer s.dstConn.Close()
+
+	checkIdleFreq := closeAfterIdle / 8
+	checkIdleTicker := time.NewTicker(checkIdleFreq)
+	defer checkIdleTicker.Stop()
+
+	activeAt := time.Now()
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case <-s.doneChan:
+			return
+		case <-checkIdleTicker.C:
+			// The session is considered inactive if current time is after (last active time + allowed idle time)
+			if time.Now().After(activeAt.Add(closeAfterIdle)) {
+				return
+			}
+		case activeAt = <-s.activeAtChan: // Update last active time
+		}
+	}
+}
+
+func (s *Session) dstToTransport(buffer []byte) error {
+	n, err := s.dstConn.Read(buffer)
+	s.markActive()
+	if n > 0 {
+		if err := s.transport.SendTo(s.id, buffer[:n]); err != nil {
+			return err
+		}
+	}
+	return err
+}
+
+func (s *Session) transportToDst(payload []byte) (int, error) {
+	s.markActive()
 	return s.dstConn.Write(payload)
 }
 
+// Sends the last active time to the idle checker loop without blocking. activeAtChan will only be full when there
+// are many concurrent read/write. It is fine to lose some precision
+func (s *Session) markActive() {
+	select {
+	case s.activeAtChan <- time.Now():
+	default:
+	}
+}
+
 func (s *Session) close() {
 	close(s.doneChan)
 }