RSDK-10204 - Camera API server GetImage data race (viamrobotics#4843)

Author: cheukt

components/camera/server.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"image"
+	"sync"
 
 	"github.com/pkg/errors"
 	"go.opencensus.io/trace"
@@ -22,9 +23,11 @@ import (
 // serviceServer implements the CameraService from camera.proto.
 type serviceServer struct {
 	pb.UnimplementedCameraServiceServer
-	coll     resource.APIResourceCollection[Camera]
-	imgTypes map[string]ImageType
-	logger   logging.Logger
+	coll resource.APIResourceCollection[Camera]
+
+	imgTypesMu sync.RWMutex
+	imgTypes   map[string]ImageType
+	logger     logging.Logger
 }
 
 // NewRPCServiceServer constructs an camera gRPC service server.
@@ -54,16 +57,22 @@ func (s *serviceServer) GetImage(
 
 	// Determine the mimeType we should try to use based on camera properties
 	if req.MimeType == "" {
-		if _, ok := s.imgTypes[req.Name]; !ok {
+		s.imgTypesMu.RLock()
+		imgType, ok := s.imgTypes[req.Name]
+		s.imgTypesMu.RUnlock()
+		if !ok {
 			props, err := cam.Properties(ctx)
 			if err != nil {
 				s.logger.CWarnf(ctx, "camera properties not found for %s, assuming color images: %v", req.Name, err)
-				s.imgTypes[req.Name] = ColorStream
+				imgType = ColorStream
 			} else {
-				s.imgTypes[req.Name] = props.ImageType
+				imgType = props.ImageType
 			}
+			s.imgTypesMu.Lock()
+			s.imgTypes[req.Name] = imgType
+			s.imgTypesMu.Unlock()
 		}
-		switch s.imgTypes[req.Name] {
+		switch imgType {
 		case ColorStream, UnspecifiedStream:
 			req.MimeType = utils.MimeTypeJPEG
 		case DepthStream:

components/camera/server_test.go

@@ -7,6 +7,7 @@ import (
 	"image"
 	"image/jpeg"
 	"image/png"
+	"sync"
 	"testing"
 	"time"
 
@@ -478,3 +479,25 @@ func TestServer(t *testing.T) {
 		test.That(t, err.Error(), test.ShouldContainSubstring, errGetImageFailed.Error())
 	})
 }
+
+func TestGetImageRace(t *testing.T) {
+	cameraServer, injectCamera, _, _, err := newServer()
+	test.That(t, err, test.ShouldBeNil)
+
+	injectCamera.PropertiesFunc = func(ctx context.Context) (camera.Properties, error) { return camera.Properties{}, nil }
+	injectCamera.ImageFunc = func(ctx context.Context, mimeType string, extra map[string]interface{}) ([]byte, camera.ImageMetadata, error) {
+		return []byte{1, 2}, camera.ImageMetadata{}, nil
+	}
+	var wg sync.WaitGroup
+
+	getImg := func() {
+		defer wg.Done()
+		_, retErr := cameraServer.GetImage(context.Background(), &pb.GetImageRequest{Name: testCameraName})
+		test.That(t, retErr, test.ShouldBeNil)
+	}
+	for range 2 {
+		wg.Add(1)
+		go getImg()
+	}
+	wg.Wait()
+}