TUN-7351: Add streaming logs session ping and timeout

Sends a ping every 15 seconds to keep the session alive even if no
protocol messages are being propagated. Additionally, sets a hard
timeout of 5 minutes when not actively streaming logs to drop the
connection.
This commit is contained in:
Devin Carr 2023-04-06 16:00:19 -07:00
parent 3fd571063e
commit 8d87d4facd
1 changed files with 22 additions and 1 deletions

View File

@ -5,6 +5,7 @@ import (
"net/http" "net/http"
"sync" "sync"
"sync/atomic" "sync/atomic"
"time"
"github.com/go-chi/chi/v5" "github.com/go-chi/chi/v5"
"github.com/rs/zerolog" "github.com/rs/zerolog"
@ -19,6 +20,9 @@ const (
// value will return this error to incoming requests. // value will return this error to incoming requests.
StatusSessionLimitExceeded websocket.StatusCode = 4002 StatusSessionLimitExceeded websocket.StatusCode = 4002
reasonSessionLimitExceeded = "limit exceeded for streaming sessions" reasonSessionLimitExceeded = "limit exceeded for streaming sessions"
StatusIdleLimitExceeded websocket.StatusCode = 4003
reasonIdleLimitExceeded = "session was idle for too long"
) )
type ManagementService struct { type ManagementService struct {
@ -147,10 +151,20 @@ func (m *ManagementService) logs(w http.ResponseWriter, r *http.Request) {
} }
// Make sure the connection is closed if other go routines fail to close the connection after completing. // Make sure the connection is closed if other go routines fail to close the connection after completing.
defer c.Close(websocket.StatusInternalError, "") defer c.Close(websocket.StatusInternalError, "")
ctx := r.Context() ctx, cancel := context.WithCancel(r.Context())
defer cancel()
events := make(chan *ClientEvent) events := make(chan *ClientEvent)
go m.readEvents(c, ctx, events) go m.readEvents(c, ctx, events)
// Send a heartbeat ping to hold the connection open even if not streaming.
ping := time.NewTicker(15 * time.Second)
defer ping.Stop()
// Close the connection if no operation has occurred after the idle timeout.
idleTimeout := 5 * time.Minute
idle := time.NewTimer(idleTimeout)
defer idle.Stop()
for { for {
select { select {
case <-ctx.Done(): case <-ctx.Done():
@ -160,9 +174,11 @@ func (m *ManagementService) logs(w http.ResponseWriter, r *http.Request) {
case event := <-events: case event := <-events:
switch event.Type { switch event.Type {
case StartStreaming: case StartStreaming:
idle.Stop()
m.startStreaming(c, ctx, event) m.startStreaming(c, ctx, event)
continue continue
case StopStreaming: case StopStreaming:
idle.Reset(idleTimeout)
// TODO: limit StopStreaming to only halt streaming for clients that are already streaming // TODO: limit StopStreaming to only halt streaming for clients that are already streaming
m.streaming.Store(false) m.streaming.Store(false)
case UnknownClientEventType: case UnknownClientEventType:
@ -176,6 +192,11 @@ func (m *ManagementService) logs(w http.ResponseWriter, r *http.Request) {
} }
return return
} }
case <-ping.C:
go c.Ping(ctx)
case <-idle.C:
c.Close(StatusIdleLimitExceeded, reasonIdleLimitExceeded)
return
} }
} }
} }