add initial poc

.gitignore

@@ -65,4 +65,7 @@ e2e/.bin
 e2e/.env.e2e
 
 # IDE
-.cursor
+.cursor
+
+# Module build output
+modules/*/bin

buf.lock

@@ -5,5 +5,5 @@ deps:
     commit: 62f35d8aed1149c291d606d958a7ce32
     digest: b5:d66bf04adc77a0870bdc9328aaf887c7188a36fb02b83a480dc45ef9dc031b4d39fc6e9dc6435120ccf4fe5bfd5c6cb6592533c6c316595571f9a31420ab47fe
   - name: buf.build/viamrobotics/api
-    commit: 17805817b2b2439aa3a600f5b2ccf3df
-    digest: b5:b40d821129b8a6a7561e1c455a2b8da8fb243da83970ad6566fc23a8652ff586f92d52c1ebb0811738a24bdf7e689c1ea5920e01e71448ba64fc1a9a2667addc
+    commit: 998407192f184de9a96c94adb8513244
+    digest: b5:1a0ed5e0080de164d97a288a8626227ee4f4bb61d6748de99063852082c1f5573c7922fe8cd0aadbb3b1bfe1b41d9e82e2ce9adb732462cf097a050fef294927

draw/draw_service.go

@@ -318,6 +318,9 @@ func (svc *DrawService) addEntityLocked(msg *drawv1.AddEntityRequest) (uuid.UUID
 			return id, connect.NewError(connect.CodeInvalidArgument, errors.New("transform is required"))
 		}
 		id = resolveEntityUUID(e.Transform.GetUuid())
+		// Stamped before storing: a server-assigned UUID is only returned to the caller, so
+		// without this the stored entity and every stream event carry an empty one.
+		e.Transform.Uuid = id[:]
 		_, exists := svc.entities[id]
 		changeType := addedOrUpdated(exists)
 		svc.entities[id] = storedEntity{kind: entityKindTransform, transform: e.Transform}
@@ -330,6 +333,7 @@ func (svc *DrawService) addEntityLocked(msg *drawv1.AddEntityRequest) (uuid.UUID
 			return id, connect.NewError(connect.CodeInvalidArgument, errors.New("drawing is required"))
 		}
 		id = resolveEntityUUID(e.Drawing.GetUuid())
+		e.Drawing.Uuid = id[:]
 		_, exists := svc.entities[id]
 		changeType := addedOrUpdated(exists)
 		svc.entities[id] = storedEntity{kind: entityKindDrawing, drawing: e.Drawing}
@@ -1207,72 +1211,27 @@ func (svc *DrawService) StreamEntityChanges(
 	_ *connect.Request[drawv1.StreamEntityChangesRequest],
 	stream *connect.ServerStream[drawv1.StreamEntityChangesResponse],
 ) error {
-	svc.mu.Lock()
-	subID, sub := svc.addEntitySub()
-
-	replay := make([]*drawv1.StreamEntityChangesResponse, 0, len(svc.entities))
-	for id, entity := range svc.entities {
-		switch entity.kind {
-		case entityKindTransform:
-			replay = append(replay, &drawv1.StreamEntityChangesResponse{
-				ChangeType: drawv1.EntityChangeType_ENTITY_CHANGE_TYPE_ADDED,
-				Entity:     &drawv1.StreamEntityChangesResponse_Transform{Transform: entity.transform},
-			})
-		case entityKindDrawing:
-			if chunked, ok := svc.chunked[id]; ok {
-				if msg := svc.buildChunkedReplayMsg(chunked); msg != nil {
-					replay = append(replay, msg)
-				}
-			} else {
-				replay = append(replay, &drawv1.StreamEntityChangesResponse{
-					ChangeType: drawv1.EntityChangeType_ENTITY_CHANGE_TYPE_ADDED,
-					Entity:     &drawv1.StreamEntityChangesResponse_Drawing{Drawing: entity.drawing},
-				})
-			}
-		}
-	}
-	svc.mu.Unlock()
-
-	// Registered before the replay loop: a client that disconnects mid-replay
-	// would otherwise leave its channel in svc.entitySubs with nobody draining
-	// it, filling up and logging a drop on every subsequent mutation forever.
-	defer func() {
-		svc.mu.Lock()
-		svc.removeEntitySub(subID)
-		svc.mu.Unlock()
-	}()
-
-	for _, msg := range replay {
-		if err := stream.Send(msg); err != nil {
-			return err
-		}
-	}
+	// Registered before the replay is sent: a client that disconnects mid-replay would
+	// otherwise leave its channel in svc.entitySubs with nobody draining it, filling up and
+	// logging a drop on every subsequent mutation forever.
+	subscription := svc.SubscribeEntities()
+	defer subscription.Close()
 
 	for {
-		msgs, overflow, closed := sub.take()
-		if overflow {
-			log.Printf("draw: subscriber %d fell more than %d changes behind; closing stream for resync", subID, maxPendingEntityChanges)
-			return connect.NewError(
-				connect.CodeResourceExhausted,
-				errors.New("entity change queue overflow; reconnect for a fresh snapshot"),
-			)
+		msgs, err := subscription.Next(ctx)
+		switch {
+		case errors.Is(err, ErrSubscriberOverflow):
+			log.Printf("draw: subscriber %d fell more than %d changes behind; closing stream for resync", subscription.id, maxPendingEntityChanges)
+			return connect.NewError(connect.CodeResourceExhausted, ErrSubscriberOverflow)
+		case err != nil:
+			return nil
 		}
+
 		for _, msg := range msgs {
 			if err := stream.Send(msg); err != nil {
 				return err
 			}
 		}
-		if len(msgs) > 0 {
-			continue
-		}
-		if closed {
-			return nil
-		}
-		select {
-		case <-ctx.Done():
-			return nil
-		case <-sub.notify:
-		}
 	}
 }
 

draw/draw_service_test.go

@@ -128,6 +128,39 @@ func TestDrawService_AddEntity(t *testing.T) {
 		test.That(t, resp.Msg.GetUuid(), test.ShouldHaveLength, 16)
 	})
 
+	t.Run("ServerAssignedUUIDReachesTheStream", func(t *testing.T) {
+		svc := NewDrawService(t.TempDir())
+		client := newTestServer(t, svc)
+
+		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+		defer cancel()
+
+		type streamResult struct {
+			stream *connect.ServerStreamForClient[drawv1.StreamEntityChangesResponse]
+			err    error
+		}
+		sCh := make(chan streamResult, 1)
+		go func() {
+			s, err := client.StreamEntityChanges(ctx, connect.NewRequest(&drawv1.StreamEntityChangesRequest{}))
+			sCh <- streamResult{s, err}
+		}()
+
+		waitForEntitySubs(t, svc, 1)
+
+		// No UUID supplied, so the service mints one. A consumer keying off the stream can only
+		// track the entity if that UUID is stamped onto the broadcast entity too.
+		resp, err := client.AddEntity(context.Background(), connect.NewRequest(&drawv1.AddEntityRequest{
+			Entity: &drawv1.AddEntityRequest_Drawing{Drawing: sampleDrawing("drawing-1")},
+		}))
+		test.That(t, err, test.ShouldBeNil)
+
+		sr := <-sCh
+		test.That(t, sr.err, test.ShouldBeNil)
+
+		test.That(t, sr.stream.Receive(), test.ShouldBeTrue)
+		test.That(t, sr.stream.Msg().GetDrawing().GetUuid(), test.ShouldResemble, resp.Msg.GetUuid())
+	})
+
 	t.Run("AddMultipleEntitiesReturnsUniqueUUIDs", func(t *testing.T) {
 		svc := NewDrawService(t.TempDir())
 		client := newTestServer(t, svc)

draw/entity_subscription.go

@@ -0,0 +1,101 @@
+package draw
+
+import (
+	"context"
+	"errors"
+	"io"
+	"sync"
+
+	drawv1 "github.com/viam-labs/motion-tools/draw/v1"
+)
+
+// ErrSubscriberOverflow reports that a subscription fell too far behind. The stream is
+// incremental, so the backlog cannot be salvaged; resubscribe for a fresh replay.
+var ErrSubscriberOverflow = errors.New("entity change queue overflow; reconnect for a fresh snapshot")
+
+// EntitySubscription is a transport-independent view of the entity change stream, so a consumer
+// outside this package can follow the scene without going through Connect-RPC.
+type EntitySubscription struct {
+	svc    *DrawService
+	id     uint64
+	sub    *entitySubscriber
+	replay []*drawv1.StreamEntityChangesResponse
+
+	closeOnce sync.Once
+}
+
+// SubscribeEntities registers a subscriber against the current world. Callers must Close it: an
+// abandoned subscription queues changes nobody drains until it overflows.
+func (svc *DrawService) SubscribeEntities() *EntitySubscription {
+	// One lock covers the snapshot and the registration, so no change can slip between them.
+	svc.mu.Lock()
+	defer svc.mu.Unlock()
+
+	id, sub := svc.addEntitySub()
+
+	replay := make([]*drawv1.StreamEntityChangesResponse, 0, len(svc.entities))
+	for entityID, entity := range svc.entities {
+		switch entity.kind {
+		case entityKindTransform:
+			replay = append(replay, &drawv1.StreamEntityChangesResponse{
+				ChangeType: drawv1.EntityChangeType_ENTITY_CHANGE_TYPE_ADDED,
+				Entity:     &drawv1.StreamEntityChangesResponse_Transform{Transform: entity.transform},
+			})
+		case entityKindDrawing:
+			if chunked, ok := svc.chunked[entityID]; ok {
+				if msg := svc.buildChunkedReplayMsg(chunked); msg != nil {
+					replay = append(replay, msg)
+				}
+			} else {
+				replay = append(replay, &drawv1.StreamEntityChangesResponse{
+					ChangeType: drawv1.EntityChangeType_ENTITY_CHANGE_TYPE_ADDED,
+					Entity:     &drawv1.StreamEntityChangesResponse_Drawing{Drawing: entity.drawing},
+				})
+			}
+		}
+	}
+
+	return &EntitySubscription{svc: svc, id: id, sub: sub, replay: replay}
+}
+
+// Next blocks for the next batch of changes: the world replayed as ADDED events on the first
+// call, live changes after.
+//
+// Returns io.EOF when closed or cancelled, ErrSubscriberOverflow when the consumer fell behind.
+func (subscription *EntitySubscription) Next(
+	ctx context.Context,
+) ([]*drawv1.StreamEntityChangesResponse, error) {
+	if len(subscription.replay) > 0 {
+		replay := subscription.replay
+		subscription.replay = nil
+		return replay, nil
+	}
+
+	for {
+		msgs, overflow, closed := subscription.sub.take()
+		if overflow {
+			return nil, ErrSubscriberOverflow
+		}
+		if len(msgs) > 0 {
+			return msgs, nil
+		}
+		if closed {
+			return nil, io.EOF
+		}
+
+		select {
+		case <-ctx.Done():
+			return nil, io.EOF
+		case <-subscription.sub.notify:
+		}
+	}
+}
+
+// Close unregisters the subscription. Safe to call more than once.
+func (subscription *EntitySubscription) Close() {
+	subscription.closeOnce.Do(func() {
+		subscription.svc.mu.Lock()
+		subscription.svc.removeEntitySub(subscription.id)
+		subscription.svc.mu.Unlock()
+	})
+}