55 "net/http"
66 "sync"
77 "sync/atomic"
8+ "time"
89
910 "github.com/go-chi/chi/v5"
1011 "github.com/rs/zerolog"
@@ -19,6 +20,9 @@ const (
1920 // value will return this error to incoming requests.
2021 StatusSessionLimitExceeded websocket.StatusCode = 4002
2122 reasonSessionLimitExceeded = "limit exceeded for streaming sessions"
23+
24+ StatusIdleLimitExceeded websocket.StatusCode = 4003
25+ reasonIdleLimitExceeded = "session was idle for too long"
2226)
2327
2428type ManagementService struct {
@@ -147,10 +151,20 @@ func (m *ManagementService) logs(w http.ResponseWriter, r *http.Request) {
147151 }
148152 // Make sure the connection is closed if other go routines fail to close the connection after completing.
149153 defer c .Close (websocket .StatusInternalError , "" )
150- ctx := r .Context ()
154+ ctx , cancel := context .WithCancel (r .Context ())
155+ defer cancel ()
151156 events := make (chan * ClientEvent )
152157 go m .readEvents (c , ctx , events )
153158
159+ // Send a heartbeat ping to hold the connection open even if not streaming.
160+ ping := time .NewTicker (15 * time .Second )
161+ defer ping .Stop ()
162+
163+ // Close the connection if no operation has occurred after the idle timeout.
164+ idleTimeout := 5 * time .Minute
165+ idle := time .NewTimer (idleTimeout )
166+ defer idle .Stop ()
167+
154168 for {
155169 select {
156170 case <- ctx .Done ():
@@ -160,9 +174,11 @@ func (m *ManagementService) logs(w http.ResponseWriter, r *http.Request) {
160174 case event := <- events :
161175 switch event .Type {
162176 case StartStreaming :
177+ idle .Stop ()
163178 m .startStreaming (c , ctx , event )
164179 continue
165180 case StopStreaming :
181+ idle .Reset (idleTimeout )
166182 // TODO: limit StopStreaming to only halt streaming for clients that are already streaming
167183 m .streaming .Store (false )
168184 case UnknownClientEventType :
@@ -176,6 +192,11 @@ func (m *ManagementService) logs(w http.ResponseWriter, r *http.Request) {
176192 }
177193 return
178194 }
195+ case <- ping .C :
196+ go c .Ping (ctx )
197+ case <- idle .C :
198+ c .Close (StatusIdleLimitExceeded , reasonIdleLimitExceeded )
199+ return
179200 }
180201 }
181202}
0 commit comments