RSDK-6340 Allow streams to restart after disconnect (viamrobotics#3467)

Author: maximpertsov

components/audioinput/client.go

@@ -42,6 +42,9 @@ func NewClientFromConn(
 	name resource.Name,
 	logger logging.Logger,
 ) (AudioInput, error) {
+	// TODO(RSDK-6340): This client might still try to create audio streams after this
+	// context is canceled. These subsequent audio streams will not work. To fix this,
+	// use a channel instead of a context like we do in `component/audioinput/client.go`
 	cancelCtx, cancel := context.WithCancel(context.Background())
 	c := pb.NewAudioInputServiceClient(conn)
 	return &client{

components/camera/client.go

@@ -31,15 +31,13 @@ import (
 type client struct {
 	resource.Named
 	resource.TriviallyReconfigurable
-	resource.TriviallyCloseable
 	mu                      sync.Mutex
 	name                    string
 	conn                    rpc.ClientConn
 	client                  pb.CameraServiceClient
 	logger                  logging.Logger
 	activeBackgroundWorkers sync.WaitGroup
-	cancelCtx               context.Context
-	cancel                  func()
+	healthyClientCh         chan struct{}
 }
 
 // NewClientFromConn constructs a new Client from connection passed in.
@@ -50,16 +48,13 @@ func NewClientFromConn(
 	name resource.Name,
 	logger logging.Logger,
 ) (Camera, error) {
-	cancelCtx, cancel := context.WithCancel(context.Background())
 	c := pb.NewCameraServiceClient(conn)
 	return &client{
-		Named:     name.PrependRemote(remoteName).AsNamed(),
-		name:      name.ShortName(),
-		conn:      conn,
-		client:    c,
-		logger:    logger,
-		cancelCtx: cancelCtx,
-		cancel:    cancel,
+		Named:  name.PrependRemote(remoteName).AsNamed(),
+		name:   name.ShortName(),
+		conn:   conn,
+		client: c,
+		logger: logger,
 	}, nil
 }
 
@@ -121,17 +116,40 @@ func (c *client) Stream(
 ) (gostream.VideoStream, error) {
 	ctx, span := trace.StartSpan(ctx, "camera::client::Stream")
 
-	cancelCtxWithMIME := gostream.WithMIMETypeHint(c.cancelCtx, gostream.MIMETypeHint(ctx, ""))
-	streamCtx, stream, frameCh := gostream.NewMediaStreamForChannel[image.Image](cancelCtxWithMIME)
-
+	// RSDK-6340: The resource manager closes remote resources when the underlying
+	// connection goes bad. However, when the connection is re-established, the client
+	// objects these resources represent are not re-initialized/marked "healthy".
+	// `healthyClientCh` helps track these transitions between healthy and unhealthy
+	// states.
+	//
+	// When a new `client.Stream()` is created we will either use the existing
+	// `healthyClientCh` or create a new one.
+	//
+	// The goroutine a `Stream()` method spins off will listen to its version of the
+	// `healthyClientCh` to be notified when the connection has died so it can gracefully
+	// terminate.
+	//
+	// When a connection becomes unhealthy, the resource manager will call `Close` on the
+	// camera client object. Closing the client will:
+	// 1. close its `client.healthyClientCh` channel
+	// 2. wait for existing "stream" goroutines to drain
+	// 3. nil out the `client.healthyClientCh` member variable
+	//
+	// New streams concurrent with closing cannot start until this drain completes. There
+	// will never be stream goroutines from the old "generation" running concurrently
+	// with those from the new "generation".
 	c.mu.Lock()
-	if err := c.cancelCtx.Err(); err != nil {
-		c.mu.Unlock()
-		return nil, err
+	if c.healthyClientCh == nil {
+		c.healthyClientCh = make(chan struct{})
 	}
-	c.activeBackgroundWorkers.Add(1)
+	healthyClientCh := c.healthyClientCh
 	c.mu.Unlock()
 
+	ctxWithMIME := gostream.WithMIMETypeHint(context.Background(), gostream.MIMETypeHint(ctx, ""))
+	streamCtx, stream, frameCh := gostream.NewMediaStreamForChannel[image.Image](ctxWithMIME)
+
+	c.activeBackgroundWorkers.Add(1)
+
 	goutils.PanicCapturingGo(func() {
 		streamCtx = trace.NewContext(streamCtx, span)
 		defer span.End()
@@ -154,6 +172,11 @@ func (c *client) Stream(
 			select {
 			case <-streamCtx.Done():
 				return
+			case <-healthyClientCh:
+				if err := stream.Close(ctxWithMIME); err != nil {
+					c.logger.Warn("error closing stream", err)
+				}
+				return
 			case frameCh <- gostream.MediaReleasePairWithError[image.Image]{
 				Media:   frame,
 				Release: release,
@@ -290,10 +313,18 @@ func (c *client) DoCommand(ctx context.Context, cmd map[string]interface{}) (map
 	return protoutils.DoFromResourceClient(ctx, c.client, c.name, cmd)
 }
 
+// TODO(RSDK-6433): This method can be called more than once during a client's lifecycle.
+// For example, consider a case where a remote camera goes offline and then back online.
+// We will call `Close` on the camera client when we detect the disconnection to remove
+// active streams but then reuse the client when the connection is re-established.
 func (c *client) Close(ctx context.Context) error {
 	c.mu.Lock()
-	c.cancel()
-	c.mu.Unlock()
+	defer c.mu.Unlock()
+
+	if c.healthyClientCh != nil {
+		close(c.healthyClientCh)
+	}
 	c.activeBackgroundWorkers.Wait()
+	c.healthyClientCh = nil
 	return nil
 }

components/camera/client_test.go

@@ -551,3 +551,75 @@ func TestClientWithInterceptor(t *testing.T) {
 
 	test.That(t, conn.Close(), test.ShouldBeNil)
 }
+
+func TestClientStreamAfterClose(t *testing.T) {
+	// Set up gRPC server
+	logger := logging.NewTestLogger(t)
+	listener, err := net.Listen("tcp", "localhost:0")
+	test.That(t, err, test.ShouldBeNil)
+	rpcServer, err := rpc.NewServer(logger.AsZap(), rpc.WithUnauthenticated())
+	test.That(t, err, test.ShouldBeNil)
+
+	// Set up camera that can stream images
+	img := image.NewNRGBA(image.Rect(0, 0, 4, 4))
+	injectCamera := &inject.Camera{}
+	injectCamera.PropertiesFunc = func(ctx context.Context) (camera.Properties, error) {
+		return camera.Properties{}, nil
+	}
+	injectCamera.StreamFunc = func(ctx context.Context, errHandlers ...gostream.ErrorHandler) (gostream.VideoStream, error) {
+		return gostream.NewEmbeddedVideoStreamFromReader(gostream.VideoReaderFunc(func(ctx context.Context) (image.Image, func(), error) {
+			return img, func() {}, nil
+		})), nil
+	}
+
+	// Register CameraService API in our gRPC server.
+	resources := map[resource.Name]camera.Camera{
+		camera.Named(testCameraName): injectCamera,
+	}
+	cameraSvc, err := resource.NewAPIResourceCollection(camera.API, resources)
+	test.That(t, err, test.ShouldBeNil)
+	resourceAPI, ok, err := resource.LookupAPIRegistration[camera.Camera](camera.API)
+	test.That(t, err, test.ShouldBeNil)
+	test.That(t, ok, test.ShouldBeTrue)
+	test.That(t, resourceAPI.RegisterRPCService(context.Background(), rpcServer, cameraSvc), test.ShouldBeNil)
+
+	// Start serving requests.
+	go rpcServer.Serve(listener)
+	defer rpcServer.Stop()
+
+	// Make client connection
+	conn, err := viamgrpc.Dial(context.Background(), listener.Addr().String(), logger)
+	test.That(t, err, test.ShouldBeNil)
+	client, err := camera.NewClientFromConn(context.Background(), conn, "", camera.Named(testCameraName), logger)
+	test.That(t, err, test.ShouldBeNil)
+
+	// Get a stream
+	stream, err := client.Stream(context.Background())
+	test.That(t, stream, test.ShouldNotBeNil)
+	test.That(t, err, test.ShouldBeNil)
+
+	// Read from stream
+	media, _, err := stream.Next(context.Background())
+	test.That(t, media, test.ShouldNotBeNil)
+	test.That(t, err, test.ShouldBeNil)
+
+	// Close client and read from stream
+	test.That(t, client.Close(context.Background()), test.ShouldBeNil)
+	media, _, err = stream.Next(context.Background())
+	test.That(t, media, test.ShouldBeNil)
+	test.That(t, err.Error(), test.ShouldContainSubstring, "context canceled")
+
+	// Get a new stream
+	stream, err = client.Stream(context.Background())
+	test.That(t, stream, test.ShouldNotBeNil)
+	test.That(t, err, test.ShouldBeNil)
+
+	// Read from the new stream
+	media, _, err = stream.Next(context.Background())
+	test.That(t, media, test.ShouldNotBeNil)
+	test.That(t, err, test.ShouldBeNil)
+
+	// Close client and connection
+	test.That(t, client.Close(context.Background()), test.ShouldBeNil)
+	test.That(t, conn.Close(), test.ShouldBeNil)
+}

components/camera/fake/camera.go

@@ -6,6 +6,7 @@ import (
 	"image"
 	"image/color"
 	"math"
+	"time"
 
 	"github.com/golang/geo/r3"
 	"github.com/pkg/errors"
@@ -57,11 +58,12 @@ func NewCamera(
 	}
 	resModel, width, height := fakeModel(newConf.Width, newConf.Height)
 	cam := &Camera{
-		Named:  conf.ResourceName().AsNamed(),
-		Model:  resModel,
-		Width:  width,
-		Height: height,
-		logger: logger,
+		Named:    conf.ResourceName().AsNamed(),
+		Model:    resModel,
+		Width:    width,
+		Height:   height,
+		Animated: newConf.Animated,
+		logger:   logger,
 	}
 	src, err := camera.NewVideoSourceFromReader(ctx, cam, resModel, camera.ColorStream)
 	if err != nil {
@@ -72,8 +74,9 @@ func NewCamera(
 
 // Config are the attributes of the fake camera config.
 type Config struct {
-	Width  int `json:"width,omitempty"`
-	Height int `json:"height,omitempty"`
+	Width    int  `json:"width,omitempty"`
+	Height   int  `json:"height,omitempty"`
+	Animated bool `json:"animated,omitempty"`
 }
 
 // Validate checks that the config attributes are valid for a fake camera.
@@ -167,6 +170,7 @@ type Camera struct {
 	Model           *transform.PinholeCameraModel
 	Width           int
 	Height          int
+	Animated        bool
 	cacheImage      *image.RGBA
 	cachePointCloud pointcloud.PointCloud
 	logger          logging.Logger
@@ -183,16 +187,27 @@ func (c *Camera) Read(ctx context.Context) (image.Image, func(), error) {
 
 	totalDist := math.Sqrt(math.Pow(0-width, 2) + math.Pow(0-height, 2))
 
+	tick := time.Now().UnixMilli() / 20
 	var x, y float64
 	for x = 0; x < width; x++ {
 		for y = 0; y < height; y++ {
 			dist := math.Sqrt(math.Pow(0-x, 2) + math.Pow(0-y, 2))
 			dist /= totalDist
 			thisColor := color.RGBA{uint8(255 - (255 * dist)), uint8(255 - (255 * dist)), uint8(0 + (255 * dist)), 255}
-			img.Set(int(x), int(y), thisColor)
+
+			var px, py int
+			if c.Animated {
+				px = int(int64(x)+tick) % int(width)
+				py = int(y)
+			} else {
+				px, py = int(x), int(y)
+			}
+			img.Set(px, py, thisColor)
 		}
 	}
-	c.cacheImage = img
+	if !c.Animated {
+		c.cacheImage = img
+	}
 	return rimage.ConvertImage(img), func() {}, nil
 }
 

gostream/media.go

@@ -417,6 +417,13 @@ func (ms *mediaStreamFromChannel[T]) Next(ctx context.Context) (T, func(), error
 	defer span.End()
 
 	var zero T
+	if ms.cancelCtx.Err() != nil {
+		return zero, nil, ms.cancelCtx.Err()
+	}
+	if ctx.Err() != nil {
+		return zero, nil, ctx.Err()
+	}
+
 	select {
 	case <-ms.cancelCtx.Done():
 		return zero, nil, ms.cancelCtx.Err()

resource/resource.go

@@ -164,7 +164,7 @@ func (t TriviallyReconfigurable) Reconfigure(ctx context.Context, deps Dependenc
 
 // TriviallyCloseable is to be embedded by any resource that does not care about
 // handling Closes. When is used, it is assumed that the resource does not need
-// to return errors when furture non-Close methods are called.
+// to return errors when future non-Close methods are called.
 type TriviallyCloseable struct{}
 
 // Close always returns no error.