Simplify session and WHEP lifecycle cleanup callbacks

Author: Sean-Der

internal/webrtc/sessions/manager/manager.go

@@ -1,7 +1,6 @@
 package manager
 
 import (
-	"context"
 	"log"
 	"maps"
 	"time"
@@ -22,7 +21,6 @@ func (m *SessionManager) Setup() {
 // Add new session
 func (m *SessionManager) addSession(profile authorization.PublicProfile) (s *session.Session, err error) {
 	log.Println("SessionManager.AddWHIPSession")
-	activeContext, activeContextCancel := context.WithCancel(context.Background())
 
 	s = &session.Session{
 
@@ -31,27 +29,19 @@ func (m *SessionManager) addSession(profile authorization.PublicProfile) (s *ses
 		MOTD:        profile.MOTD,
 		StreamStart: time.Now(),
 
-		ActiveContext:       activeContext,
-		ActiveContextCancel: activeContextCancel,
-
 		WHEPSessions: map[string]*whep.WHEPSession{},
 	}
-
-	s.HasHost.Store(true)
-	m.sessionsLock.Lock()
-	m.sessions[profile.StreamKey] = s
-	m.sessionsLock.Unlock()
-
-	go func() {
-		<-activeContext.Done()
+	s.SetOnClose(func() {
 		log.Println("SessionManager.Session.Done")
-
 		m.sessionsLock.Lock()
 		delete(m.sessions, profile.StreamKey)
 		m.sessionsLock.Unlock()
+	})
 
-	}()
-
+	s.HasHost.Store(true)
+	m.sessionsLock.Lock()
+	m.sessions[profile.StreamKey] = s
+	m.sessionsLock.Unlock()
 	return s, nil
 }
 

internal/webrtc/sessions/session/routines.go

@@ -23,16 +23,6 @@ import (
 //
 // }
 
-// Waits for WHEP disconnect and removes the session
-func (s *Session) handleWHEPConnection(whepSession *whep.WHEPSession) {
-	log.Println("Session.WHEPSession.Connected:", s.StreamKey)
-
-	<-whepSession.ActiveContext.Done()
-
-	log.Println("Session.WHEPSession.Disconnected:", s.StreamKey, " - ", whepSession.SessionID)
-	s.removeWHEP(whepSession.SessionID)
-}
-
 func (s *Session) handleWHEPVideoRTCPSender(whepSession *whep.WHEPSession, rtcpSender *webrtc.RTPSender) {
 	for {
 		rtcpPackets, _, rtcpErr := rtcpSender.ReadRTCP()

internal/webrtc/sessions/session/session.go

@@ -1,7 +1,6 @@
 package session
 
 import (
-	"context"
 	"fmt"
 	"log"
 
@@ -23,6 +22,10 @@ func (s *Session) UpdateStreamStatus(profile authorization.PublicProfile) {
 	s.StatusLock.Unlock()
 }
 
+func (session *Session) SetOnClose(onClose func()) {
+	session.onClose = onClose
+}
+
 // Add WHEP session to existing WHIP session
 func (s *Session) AddWHEP(whepSessionID string, peerConnection *webrtc.PeerConnection, audioTrack *codecs.TrackMultiCodec, videoTrack *codecs.TrackMultiCodec, videoRTCPSender *webrtc.RTPSender) (err error) {
 	log.Println("WHIPSessionManager.WHIPSession.AddWHEPSession")
@@ -47,8 +50,7 @@ func (s *Session) AddWHEP(whepSessionID string, peerConnection *webrtc.PeerConne
 	s.WHEPSessions[whepSessionID] = whepSession
 	s.WHEPSessionsLock.Unlock()
 	s.updateHostWHEPSessionsSnapshot()
-
-	go s.handleWHEPConnection(whepSession)
+	whepSession.SetOnClose(s.handleWHEPClose)
 	go s.handleWHEPVideoRTCPSender(whepSession, videoRTCPSender)
 
 	return nil
@@ -64,7 +66,7 @@ func (s *Session) AddHost(peerConnection *webrtc.PeerConnection) (err error) {
 			break
 		}
 
-		if host.PeerConnection.ConnectionState() != webrtc.PeerConnectionStateClosed || s.ActiveContext.Err() == nil {
+		if host.PeerConnection.ConnectionState() != webrtc.PeerConnectionStateClosed {
 			return fmt.Errorf("session already has a host")
 		}
 
@@ -73,20 +75,14 @@ func (s *Session) AddHost(peerConnection *webrtc.PeerConnection) (err error) {
 		}
 	}
 
-	activeContext, activeContextCancel := context.WithCancel(context.Background())
-
 	host := &whip.WHIPSession{
 		ID:          uuid.New().String(),
 		AudioTracks: make(map[string]*whip.AudioTrack),
 		VideoTracks: make(map[string]*whip.VideoTrack),
-
-		ActiveContext:       activeContext,
-		ActiveContextCancel: activeContextCancel,
 	}
 
 	host.AddPeerConnection(peerConnection, s.StreamKey)
 	if !s.Host.CompareAndSwap(nil, host) {
-		host.ActiveContextCancel()
 		host.RemovePeerConnection()
 		host.RemoveTracks()
 		return fmt.Errorf("session already has a host")
@@ -110,25 +106,24 @@ func (s *Session) RemoveHost() {
 	log.Println("Session.RemoveHost", s.StreamKey)
 
 	host.WHEPSessionsSnapshot.Store(make(map[string]*whep.WHEPSession))
-	host.ActiveContextCancel()
 	host.RemovePeerConnection()
 	host.RemoveTracks()
 }
 
-// Remove WHEP session from WHIP session
-// In case the WHIP session does not have a host, and no more whep sessions, it will
-// be remove from the manager.
-func (s *Session) removeWHEP(whepSessionID string) {
-	log.Println("Session.RemoveWHEPSession:", s.StreamKey, " - ", whepSessionID)
+func (s *Session) handleWHEPClose(whepSessionID string) {
+	log.Println("Session.HandleWHEPClose:", s.StreamKey, " - ", whepSessionID)
 
 	s.WHEPSessionsLock.Lock()
-	if whepSession, ok := s.WHEPSessions[whepSessionID]; ok {
-		whepSession.Close()
+	_, ok := s.WHEPSessions[whepSessionID]
+	if ok {
 		delete(s.WHEPSessions, whepSessionID)
-	} else {
-		log.Println("Session.RemoveWHEPSession.InvalidSession:", s.StreamKey, " - ", whepSessionID)
 	}
 	s.WHEPSessionsLock.Unlock()
+
+	if !ok {
+		return
+	}
+
 	s.updateHostWHEPSessionsSnapshot()
 
 	if s.isEmpty() {
@@ -138,22 +133,26 @@ func (s *Session) removeWHEP(whepSessionID string) {
 
 // Remove all Hosts and clients before closing down session
 func (s *Session) close() {
-	s.WHEPSessionsLock.Lock()
-	whepSessions := make([]*whep.WHEPSession, 0, len(s.WHEPSessions))
-	for _, whepSession := range s.WHEPSessions {
-		whepSessions = append(whepSessions, whepSession)
-	}
-	s.WHEPSessions = make(map[string]*whep.WHEPSession)
-	s.WHEPSessionsLock.Unlock()
+	s.closeOnce.Do(func() {
+		s.WHEPSessionsLock.Lock()
+		whepSessions := make([]*whep.WHEPSession, 0, len(s.WHEPSessions))
+		for _, whepSession := range s.WHEPSessions {
+			whepSessions = append(whepSessions, whepSession)
+		}
+		s.WHEPSessions = make(map[string]*whep.WHEPSession)
+		s.WHEPSessionsLock.Unlock()
 
-	for _, whepSession := range whepSessions {
-		whepSession.Close()
-	}
-	s.updateHostWHEPSessionsSnapshot()
+		for _, whepSession := range whepSessions {
+			whepSession.Close()
+		}
+		s.updateHostWHEPSessionsSnapshot()
 
-	s.RemoveHost()
+		s.RemoveHost()
 
-	s.ActiveContextCancel()
+		if s.onClose != nil {
+			s.onClose()
+		}
+	})
 }
 
 func (s *Session) Close() {

internal/webrtc/sessions/session/types.go

@@ -1,7 +1,6 @@
 package session
 
 import (
-	"context"
 	"sync"
 	"sync/atomic"
 	"time"
@@ -23,9 +22,8 @@ type Session struct {
 
 	Host atomic.Pointer[whip.WHIPSession]
 
-	// Context
-	ActiveContext       context.Context
-	ActiveContextCancel func()
+	closeOnce sync.Once
+	onClose   func()
 
 	// Protects WHEPSessions
 	WHEPSessionsLock sync.RWMutex

internal/webrtc/sessions/whep/types.go

@@ -1,7 +1,6 @@
 package whep
 
 import (
-	"context"
 	"sync"
 	"sync/atomic"
 	"time"
@@ -21,12 +20,11 @@ type (
 		IsWaitingForKeyframe atomic.Bool
 		IsSessionClosed      atomic.Bool
 
-		SSESubscribersLock  sync.RWMutex
-		SSESubscribers      map[string]sseSubscriber
-		SessionClose        sync.Once
-		ActiveContext       context.Context
-		ActiveContextCancel func()
-		pliSender           func()
+		SSESubscribersLock sync.RWMutex
+		SSESubscribers     map[string]sseSubscriber
+		SessionClose       sync.Once
+		onClose            func(string)
+		pliSender          func()
 
 		PeerConnectionLock sync.RWMutex
 		PeerConnection     *webrtc.PeerConnection

internal/webrtc/sessions/whep/whep.go

@@ -1,7 +1,6 @@
 package whep
 
 import (
-	"context"
 	"log"
 	"time"
 
@@ -21,7 +20,6 @@ func CreateNewWHEP(
 ) (w *WHEPSession) {
 	log.Println("WHEPSession.CreateNewWHEP", whepSessionID)
 
-	activeContext, activeContextCancel := context.WithCancel(context.Background())
 	w = &WHEPSession{
 		SessionID:               whepSessionID,
 		AudioTrack:              audioTrack,
@@ -30,8 +28,6 @@ func CreateNewWHEP(
 		VideoTimestamp:          5000,
 		SSESubscribers:          make(map[string]sseSubscriber),
 		PeerConnection:          peerConnection,
-		ActiveContext:           activeContext,
-		ActiveContextCancel:     activeContextCancel,
 		pliSender:               pliSender,
 		videoBitrateWindowStart: time.Now(),
 	}
@@ -74,10 +70,16 @@ func (w *WHEPSession) Close() {
 		w.SSESubscribers = make(map[string]sseSubscriber)
 		w.SSESubscribersLock.Unlock()
 
-		w.ActiveContextCancel()
+		if w.onClose != nil {
+			w.onClose(w.SessionID)
+		}
 	})
 }
 
+func (w *WHEPSession) SetOnClose(onClose func(string)) {
+	w.onClose = onClose
+}
+
 // Get the current status of the WHEP session
 func (w *WHEPSession) GetWHEPSessionStatus() (state SessionState) {
 	w.AudioLock.RLock()