jamesread 1 день назад
Родитель
Сommit
c2d9c28a59
2 измененных файлов с 89 добавлено и 9 удалено
  1. 36 9
      service/internal/api/api.go
  2. 53 0
      service/internal/api/api_test.go

+ 36 - 9
service/internal/api/api.go

@@ -1118,22 +1118,49 @@ func (api *oliveTinAPI) EventStream(ctx ctx.Context, req *connect.Request[apiv1.
 		"authenticatedUser": user.Username,
 	}).Debugf("EventStream: client connected")
 
+	// Release the slot on every exit path. removeClient is idempotent, so this
+	// is safe when a broadcast or the heartbeat goroutine already removed us.
+	defer api.removeClient(client)
+
 	go api.sendEventStreamHeartbeats(client)
 
-	// loop over client channel and send events to connectedClient
-	for msg := range client.channel {
+	api.forwardEventsToClient(ctx.Done(), client, srv)
+
+	log.Infof("EventStream: client disconnected")
+
+	return nil
+}
+
+// forwardEventsToClient sends queued events to the stream until the client goes
+// away. Without watching done, a client that disconnected was only noticed once
+// a later Send failed, which net/http's buffered writer reports several
+// heartbeats late (or not at all while the TCP peer keeps ACKing), so its slot
+// in maxEventStreamClients stayed taken.
+func (api *oliveTinAPI) forwardEventsToClient(done <-chan struct{}, client *streamingClient, srv *connect.ServerStream[apiv1.EventStreamResponse]) {
+	for {
+		msg, ok := nextEvent(done, client.channel)
+		if !ok {
+			return
+		}
+
 		log.Debugf("Sending event to client: %v", msg)
+
 		if err := srv.Send(msg); err != nil {
-			log.Errorf("Error sending event to client: %v", err)
-			// Remove disconnected client from the list
-			api.removeClient(client)
-			break
+			log.Warnf("Error sending event to client: %v", err)
+			return
 		}
 	}
+}
 
-	log.Infof("EventStream: client disconnected")
-
-	return nil
+// nextEvent waits for the next queued event. ok is false once the request is
+// done (the client went away) or the channel was closed by removeClient.
+func nextEvent(done <-chan struct{}, events <-chan *apiv1.EventStreamResponse) (*apiv1.EventStreamResponse, bool) {
+	select {
+	case <-done:
+		return nil, false
+	case msg, ok := <-events:
+		return msg, ok
+	}
 }
 
 func (api *oliveTinAPI) registerStreamingClient(client *streamingClient) error {

+ 53 - 0
service/internal/api/api_test.go

@@ -897,6 +897,59 @@ func TestRegisterStreamingClientEnforcesLimit(t *testing.T) {
 	api.removeClient(overflow)
 }
 
+// TestEventStreamReleasesSlotWhenClientDisconnects fills every EventStream slot,
+// disconnects all those clients and expects the slots to be released right
+// away, so a new client is accepted instead of being rejected with
+// "concurrent client limit reached".
+func TestEventStreamReleasesSlotWhenClientDisconnects(t *testing.T) {
+	cfg := config.DefaultConfig()
+	ex := executor.DefaultExecutor(cfg)
+	ex.RebuildActionMap()
+
+	ts, client := getNewTestServerAndClientWithExecutor(cfg, ex)
+	defer ts.Close()
+
+	server := ensureExecutorListener(ex)
+
+	cancels := make([]context.CancelFunc, 0, maxEventStreamClients)
+	for i := 0; i < maxEventStreamClients; i++ {
+		cancels = append(cancels, openEventStreamAndReadHeartbeat(t, client))
+	}
+	require.Equal(t, maxEventStreamClients, countStreamingClients(server))
+
+	for _, cancel := range cancels {
+		cancel()
+	}
+
+	// The next heartbeat is 10s away, so anything under that proves the
+	// disconnect itself released the slots, not a later failed write.
+	require.Eventually(t, func() bool {
+		return countStreamingClients(server) == 0
+	}, 2*time.Second, 10*time.Millisecond, "disconnected EventStream clients still hold their slots")
+
+	cancel := openEventStreamAndReadHeartbeat(t, client)
+	cancel()
+}
+
+func openEventStreamAndReadHeartbeat(t *testing.T, client apiv1connect.OliveTinApiServiceClient) context.CancelFunc {
+	t.Helper()
+
+	ctx, cancel := context.WithCancel(context.Background())
+	stream, err := client.EventStream(ctx, connect.NewRequest(&apiv1.EventStreamRequest{}))
+	require.NoError(t, err)
+	require.True(t, stream.Receive(), "expected an initial heartbeat, got error: %v", stream.Err())
+	require.NotNil(t, stream.Msg().GetHeartbeat())
+
+	return cancel
+}
+
+func countStreamingClients(api *oliveTinAPI) int {
+	api.streamingClientsMutex.RLock()
+	defer api.streamingClientsMutex.RUnlock()
+
+	return len(api.streamingClients)
+}
+
 func addEventStreamTestClients(t *testing.T, api *oliveTinAPI, lowUser, adminUser *authpublic.AuthenticatedUser) (*streamingClient, *streamingClient) {
 	t.Helper()
 	clientLow := &streamingClient{