clean up

pkg/config/config.go

@@ -312,28 +312,28 @@ func (l LimitConfig) CheckAttributesSize(attributes map[string]string) bool {
 	return uint32(total) <= l.MaxAttributesSize
 }
 
-func (l LimitConfig) CheckAsyncAttributesSize(asyncAttributes map[string][]byte) bool {
+func (l LimitConfig) CheckAsyncAttributesSize(asyncAttributes []*livekit.DataTrackSchemaDefinition) bool {
 	if l.MaxAsyncAttributesSize == 0 {
 		return true
 	}
 
 	total := 0
-	for k, v := range asyncAttributes {
-		total += len(k) + len(v)
+	for _, asyncAttribute := range asyncAttributes {
+		total += len(asyncAttribute.Id.Name) + len(asyncAttribute.Definition)
 	}
 	return uint32(total) <= l.MaxAsyncAttributesSize
 }
 
-func (l LimitConfig) CanAddAsyncAttribute(asyncAttributes map[string][]byte, toAddKey string, toAddValue []byte) bool {
+func (l LimitConfig) CanAddAsyncAttribute(asyncAttributes []*livekit.DataTrackSchemaDefinition, toAdd *livekit.DataTrackSchemaDefinition) bool {
 	if l.MaxAsyncAttributesSize == 0 {
 		return true
 	}
 
 	total := 0
-	for k, v := range asyncAttributes {
-		total += len(k) + len(v)
+	for _, asyncAttribute := range asyncAttributes {
+		total += len(asyncAttribute.Id.Name) + len(asyncAttribute.Definition)
 	}
-	return uint32(total+len(toAddKey)+len(toAddValue)) <= l.MaxAsyncAttributesSize
+	return uint32(total+len(toAdd.Id.Name)+len(toAdd.Definition)) <= l.MaxAsyncAttributesSize
 }
 
 // ---------------------------------

pkg/rtc/participant_async_attributes.go

@@ -22,6 +22,7 @@ import (
 
 	"github.com/livekit/protocol/livekit"
 	"github.com/livekit/protocol/logger"
+	"github.com/livekit/protocol/utils"
 )
 
 type ParticipantAsyncAttributesParams struct {
@@ -31,25 +32,25 @@ type ParticipantAsyncAttributesParams struct {
 type ParticipantAsyncAttributes struct {
 	params     ParticipantAsyncAttributesParams
 	lock       sync.Mutex
-	attributes map[string][]byte
+	attributes map[string]*livekit.DataTrackSchemaDefinition
 }
 
 func NewParticipantAsyncAttributes(params ParticipantAsyncAttributesParams) *ParticipantAsyncAttributes {
 	return &ParticipantAsyncAttributes{
 		params:     params,
-		attributes: make(map[string][]byte),
+		attributes: make(map[string]*livekit.DataTrackSchemaDefinition),
 	}
 }
 
-func (p *ParticipantAsyncAttributes) Add(id *livekit.DataTrackSchemaId, value []byte) {
+func (p *ParticipantAsyncAttributes) Add(aa *livekit.DataTrackSchemaDefinition) {
 	p.lock.Lock()
 	defer p.lock.Unlock()
 
-	if id == nil {
+	if aa.Id == nil {
 		return
 	}
 
-	p.attributes[ToParticipantAsyncAttributeKey(id)] = value
+	p.attributes[ToParticipantAsyncAttributeKey(aa.Id)] = utils.CloneProto(aa)
 }
 
 func (p *ParticipantAsyncAttributes) Delete(id *livekit.DataTrackSchemaId) {
@@ -71,22 +72,35 @@ func (p *ParticipantAsyncAttributes) Get(id *livekit.DataTrackSchemaId) *livekit
 		return nil
 	}
 
-	value, ok := p.attributes[ToParticipantAsyncAttributeKey(id)]
+	aa, ok := p.attributes[ToParticipantAsyncAttributeKey(id)]
 	if !ok {
 		return nil
 	}
 
-	return &livekit.DataTrackSchemaDefinition{
-		Id:         id,
-		Definition: value,
-	}
+	return aa
 }
 
-func (p *ParticipantAsyncAttributes) GetAll() map[string][]byte {
+func (p *ParticipantAsyncAttributes) GetAll() []*livekit.DataTrackSchemaDefinition {
 	p.lock.Lock()
 	defer p.lock.Unlock()
 
-	return p.attributes
+	all := make([]*livekit.DataTrackSchemaDefinition, 0, len(p.attributes))
+	for _, aa := range p.attributes {
+		all = append(all, utils.CloneProto(aa))
+	}
+	return all
+}
+
+func (p *ParticipantAsyncAttributes) GetAllIDs() []*livekit.DataTrackSchemaId {
+	p.lock.Lock()
+	defer p.lock.Unlock()
+
+	ids := make([]*livekit.DataTrackSchemaId, 0, len(p.attributes))
+	for _, aa := range p.attributes {
+		ids = append(ids, utils.CloneProto(aa.Id))
+	}
+
+	return ids
 }
 
 // -------------------------------
@@ -95,7 +109,7 @@ func ToParticipantAsyncAttributeKey(id *livekit.DataTrackSchemaId) string {
 	return fmt.Sprintf("%s/%d", id.Name, id.Encoding)
 }
 
-func FromParticipantAsyncAttributeKey(key string) *livekit.DataTrackSchemaId {
+func fromParticipantAsyncAttributeKey(key string) *livekit.DataTrackSchemaId {
 	parts := strings.Split(key, "/")
 	encoding, _ := strconv.Atoi(parts[1])
 	return &livekit.DataTrackSchemaId{

pkg/rtc/participant_async_attributes_handler.go

@@ -57,11 +57,7 @@ func (p *ParticipantImpl) HandleDefineDataTrackSchemaRequest(req *livekit.Define
 		return
 	}
 
-	if !p.params.LimitConfig.CanAddAsyncAttribute(
-		p.asyncAttributes.GetAll(),
-		ToParticipantAsyncAttributeKey(req.SchemaDefinition.Id),
-		req.SchemaDefinition.Definition,
-	) {
+	if !p.params.LimitConfig.CanAddAsyncAttribute(p.asyncAttributes.GetAll(), req.SchemaDefinition) {
 		p.sendRequestResponse(&livekit.RequestResponse{
 			Reason:  livekit.RequestResponse_LIMIT_EXCEEDED,
 			Message: "async attribute definition exceeds limit",
@@ -91,7 +87,7 @@ func (p *ParticipantImpl) HandleGetDataTrackSchemaRequest(req *livekit.GetDataTr
 }
 
 func (p *ParticipantImpl) AddDataTrackSchema(definition *livekit.DataTrackSchemaDefinition) {
-	p.asyncAttributes.Add(definition.Id, definition.Definition)
+	p.asyncAttributes.Add(definition)
 }
 
 func (p *ParticipantImpl) GetDataTrackSchema(id *livekit.DataTrackSchemaId) *livekit.DataTrackSchemaDefinition {
@@ -124,3 +120,7 @@ func (p *ParticipantImpl) ProcessGetDataTrackSchemaRequest(req *livekit.GetDataT
 
 	p.sendGetDataTrackSchemaResponse(asyncAttribute)
 }
+
+func (p *ParticipantImpl) GetAllAsyncAttributeIDs() []*livekit.DataTrackSchemaId {
+	return p.asyncAttributes.GetAllIDs()
+}

pkg/rtc/participant_async_attributes_handler_test.go

@@ -242,7 +242,11 @@ func TestGetDataTrackSchema(t *testing.T) {
 	}
 	require.Nil(t, p.GetDataTrackSchema(id))
 
-	p.asyncAttributes.Add(id, []byte("definition"))
+	definition := &livekit.DataTrackSchemaDefinition{
+		Id:         id,
+		Definition: []byte("definition"),
+	}
+	p.asyncAttributes.Add(definition)
 	got := p.GetDataTrackSchema(id)
 	require.NotNil(t, got)
 	require.Equal(t, id.Name, got.Id.Name)
@@ -317,3 +321,75 @@ func TestProcessGetDataTrackSchemaRequest(t *testing.T) {
 		require.Equal(t, def, response.GetDataTrackSchemaResponse.SchemaDefinition)
 	})
 }
+
+func TestGetAllAsyncAttributeIDs(t *testing.T) {
+	t.Run("returns empty slice when no schemas defined", func(t *testing.T) {
+		p := newParticipantWithAsyncAttributes(t, true, 0)
+		require.Empty(t, p.GetAllAsyncAttributeIDs())
+	})
+
+	t.Run("returns ids for all defined schemas", func(t *testing.T) {
+		p := newParticipantWithAsyncAttributes(t, true, 0)
+
+		id1 := &livekit.DataTrackSchemaId{
+			Name:     "schema-1",
+			Encoding: livekit.DataTrackSchemaEncoding_DATA_TRACK_SCHEMA_ENCODING_PROTOBUF,
+		}
+		id2 := &livekit.DataTrackSchemaId{
+			Name:     "schema-2",
+			Encoding: livekit.DataTrackSchemaEncoding_DATA_TRACK_SCHEMA_ENCODING_FLATBUFFER,
+		}
+		id3 := &livekit.DataTrackSchemaId{
+			Name:     "schema-1",
+			Encoding: livekit.DataTrackSchemaEncoding_DATA_TRACK_SCHEMA_ENCODING_JSON_SCHEMA,
+		}
+
+		p.AddDataTrackSchema(&livekit.DataTrackSchemaDefinition{Id: id1, Definition: []byte("def-1")})
+		p.AddDataTrackSchema(&livekit.DataTrackSchemaDefinition{Id: id2, Definition: []byte("def-2")})
+		p.AddDataTrackSchema(&livekit.DataTrackSchemaDefinition{Id: id3, Definition: []byte("def-3")})
+
+		ids := p.GetAllAsyncAttributeIDs()
+		require.Len(t, ids, 3)
+
+		got := make(map[string]livekit.DataTrackSchemaEncoding, len(ids))
+		for _, id := range ids {
+			got[id.Name+"|"+id.Encoding.String()] = id.Encoding
+		}
+		require.Contains(t, got, "schema-1|"+id1.Encoding.String())
+		require.Contains(t, got, "schema-2|"+id2.Encoding.String())
+		require.Contains(t, got, "schema-1|"+id3.Encoding.String())
+	})
+
+	t.Run("reflects deletes", func(t *testing.T) {
+		p := newParticipantWithAsyncAttributes(t, true, 0)
+
+		id := &livekit.DataTrackSchemaId{
+			Name:     "schema-1",
+			Encoding: livekit.DataTrackSchemaEncoding_DATA_TRACK_SCHEMA_ENCODING_PROTOBUF,
+		}
+		p.AddDataTrackSchema(&livekit.DataTrackSchemaDefinition{Id: id, Definition: []byte("def")})
+		require.Len(t, p.GetAllAsyncAttributeIDs(), 1)
+
+		p.asyncAttributes.Delete(id)
+		require.Empty(t, p.GetAllAsyncAttributeIDs())
+	})
+
+	t.Run("returned ids are cloned and do not alias internal state", func(t *testing.T) {
+		p := newParticipantWithAsyncAttributes(t, true, 0)
+
+		id := &livekit.DataTrackSchemaId{
+			Name:     "schema-1",
+			Encoding: livekit.DataTrackSchemaEncoding_DATA_TRACK_SCHEMA_ENCODING_PROTOBUF,
+		}
+		p.AddDataTrackSchema(&livekit.DataTrackSchemaDefinition{Id: id, Definition: []byte("def")})
+
+		ids := p.GetAllAsyncAttributeIDs()
+		require.Len(t, ids, 1)
+		ids[0].Name = "mutated"
+
+		// mutating the returned slice must not affect what's stored
+		again := p.GetAllAsyncAttributeIDs()
+		require.Len(t, again, 1)
+		require.Equal(t, "schema-1", again[0].Name)
+	})
+}

pkg/rtc/participant_async_attributes_test.go

@@ -39,7 +39,7 @@ func TestParticipantAsyncAttributes_AddAndGet(t *testing.T) {
 	}
 	value := []byte("definition-bytes")
 
-	a.Add(id, value)
+	a.Add(&livekit.DataTrackSchemaDefinition{Id: id, Definition: value})
 
 	got := a.Get(id)
 	require.NotNil(t, got)
@@ -55,14 +55,18 @@ func TestParticipantAsyncAttributes_AddOverwrites(t *testing.T) {
 		Name:     "schema-1",
 		Encoding: livekit.DataTrackSchemaEncoding_DATA_TRACK_SCHEMA_ENCODING_PROTOBUF,
 	}
-	a.Add(id, []byte("v1"))
-	a.Add(id, []byte("v2"))
+	a.Add(&livekit.DataTrackSchemaDefinition{Id: id, Definition: []byte("v1")})
+	a.Add(&livekit.DataTrackSchemaDefinition{Id: id, Definition: []byte("v2")})
 
 	got := a.Get(id)
 	require.NotNil(t, got)
 	require.Equal(t, []byte("v2"), got.Definition)
 
 	require.Len(t, a.GetAll(), 1)
+	allIDs := a.GetAllIDs()
+	require.Len(t, allIDs, 1)
+	require.Equal(t, id.Name, allIDs[0].Name)
+	require.Equal(t, id.Encoding, allIDs[0].Encoding)
 }
 
 func TestParticipantAsyncAttributes_DifferentEncodingsAreDistinct(t *testing.T) {
@@ -77,8 +81,8 @@ func TestParticipantAsyncAttributes_DifferentEncodingsAreDistinct(t *testing.T)
 		Encoding: livekit.DataTrackSchemaEncoding_DATA_TRACK_SCHEMA_ENCODING_JSON_SCHEMA,
 	}
 
-	a.Add(idProto, []byte("proto-def"))
-	a.Add(idJSON, []byte("json-def"))
+	a.Add(&livekit.DataTrackSchemaDefinition{Id: idProto, Definition: []byte("proto-def")})
+	a.Add(&livekit.DataTrackSchemaDefinition{Id: idJSON, Definition: []byte("json-def")})
 
 	gotProto := a.Get(idProto)
 	require.NotNil(t, gotProto)
@@ -89,6 +93,7 @@ func TestParticipantAsyncAttributes_DifferentEncodingsAreDistinct(t *testing.T)
 	require.Equal(t, []byte("json-def"), gotJSON.Definition)
 
 	require.Len(t, a.GetAll(), 2)
+	require.Len(t, a.GetAllIDs(), 2)
 }
 
 func TestParticipantAsyncAttributes_Delete(t *testing.T) {
@@ -98,7 +103,10 @@ func TestParticipantAsyncAttributes_Delete(t *testing.T) {
 		Name:     "schema-1",
 		Encoding: livekit.DataTrackSchemaEncoding_DATA_TRACK_SCHEMA_ENCODING_PROTOBUF,
 	}
-	a.Add(id, []byte("definition"))
+	a.Add(&livekit.DataTrackSchemaDefinition{
+		Id:         id,
+		Definition: []byte("definition"),
+	})
 
 	a.Delete(id)
 	require.Nil(t, a.Get(id))
@@ -113,7 +121,7 @@ func TestParticipantAsyncAttributes_NilId(t *testing.T) {
 	a := newTestAsyncAttributes()
 
 	// nil id should be silently ignored, not panic
-	a.Add(nil, []byte("definition"))
+	a.Add(&livekit.DataTrackSchemaDefinition{Definition: []byte("definition")})
 	require.Empty(t, a.GetAll())
 
 	require.Nil(t, a.Get(nil))
@@ -144,13 +152,21 @@ func TestParticipantAsyncAttributes_GetAllContents(t *testing.T) {
 		Encoding: livekit.DataTrackSchemaEncoding_DATA_TRACK_SCHEMA_ENCODING_FLATBUFFER,
 	}
 
-	a.Add(id1, []byte("def-1"))
-	a.Add(id2, []byte("def-2"))
+	a.Add(&livekit.DataTrackSchemaDefinition{Id: id1, Definition: []byte("def-1")})
+	a.Add(&livekit.DataTrackSchemaDefinition{Id: id2, Definition: []byte("def-2")})
 
 	all := a.GetAll()
 	require.Len(t, all, 2)
-	require.Equal(t, []byte("def-1"), all[ToParticipantAsyncAttributeKey(id1)])
-	require.Equal(t, []byte("def-2"), all[ToParticipantAsyncAttributeKey(id2)])
+	for _, aa := range all {
+		switch aa.Id.Name {
+		case "schema-1":
+			require.Equal(t, []byte("def-1"), aa.Definition)
+		case "schema-2":
+			require.Equal(t, []byte("def-2"), aa.Definition)
+		default:
+			require.Fail(t, "unexpected name", aa.Id.Name)
+		}
+	}
 }
 
 func TestParticipantAsyncAttributes_ConcurrentAccess(t *testing.T) {
@@ -169,7 +185,7 @@ func TestParticipantAsyncAttributes_ConcurrentAccess(t *testing.T) {
 					Name:     "schema",
 					Encoding: livekit.DataTrackSchemaEncoding(g % 8),
 				}
-				a.Add(id, []byte("v"))
+				a.Add(&livekit.DataTrackSchemaDefinition{Id: id, Definition: []byte("v")})
 				_ = a.Get(id)
 				_ = a.GetAll()
 				if i%3 == 0 {
@@ -199,7 +215,7 @@ func TestToKeyFromKey(t *testing.T) {
 	for _, id := range ids {
 		t.Run(id.Name, func(t *testing.T) {
 			key := ToParticipantAsyncAttributeKey(id)
-			roundTripped := FromParticipantAsyncAttributeKey(key)
+			roundTripped := fromParticipantAsyncAttributeKey(key)
 			require.Equal(t, id.Name, roundTripped.Name)
 			require.Equal(t, id.Encoding, roundTripped.Encoding)
 		})

pkg/rtc/types/interfaces.go

@@ -576,6 +576,8 @@ type LocalParticipant interface {
 	ClearParticipantListener()
 
 	GetNextSubscribedDataTrackHandle() uint16
+
+	GetAllAsyncAttributeIDs() []*livekit.DataTrackSchemaId
 }
 
 // ---------------------------------------------

pkg/rtc/types/typesfakes/fake_local_participant.go

@@ -221,6 +221,16 @@ type FakeLocalParticipant struct {
 	getAdaptiveStreamReturnsOnCall map[int]struct {
 		result1 bool
 	}
+	GetAllAsyncAttributeIDsStub        func() []*livekit.DataTrackSchemaId
+	getAllAsyncAttributeIDsMutex       sync.RWMutex
+	getAllAsyncAttributeIDsArgsForCall []struct {
+	}
+	getAllAsyncAttributeIDsReturns struct {
+		result1 []*livekit.DataTrackSchemaId
+	}
+	getAllAsyncAttributeIDsReturnsOnCall map[int]struct {
+		result1 []*livekit.DataTrackSchemaId
+	}
 	GetAnswerStub        func() (webrtc.SessionDescription, uint32, error)
 	getAnswerMutex       sync.RWMutex
 	getAnswerArgsForCall []struct {
@@ -2583,6 +2593,59 @@ func (fake *FakeLocalParticipant) GetAdaptiveStreamReturnsOnCall(i int, result1
 	}{result1}
 }
 
+func (fake *FakeLocalParticipant) GetAllAsyncAttributeIDs() []*livekit.DataTrackSchemaId {
+	fake.getAllAsyncAttributeIDsMutex.Lock()
+	ret, specificReturn := fake.getAllAsyncAttributeIDsReturnsOnCall[len(fake.getAllAsyncAttributeIDsArgsForCall)]
+	fake.getAllAsyncAttributeIDsArgsForCall = append(fake.getAllAsyncAttributeIDsArgsForCall, struct {
+	}{})
+	stub := fake.GetAllAsyncAttributeIDsStub
+	fakeReturns := fake.getAllAsyncAttributeIDsReturns
+	fake.recordInvocation("GetAllAsyncAttributeIDs", []interface{}{})
+	fake.getAllAsyncAttributeIDsMutex.Unlock()
+	if stub != nil {
+		return stub()
+	}
+	if specificReturn {
+		return ret.result1
+	}
+	return fakeReturns.result1
+}
+
+func (fake *FakeLocalParticipant) GetAllAsyncAttributeIDsCallCount() int {
+	fake.getAllAsyncAttributeIDsMutex.RLock()
+	defer fake.getAllAsyncAttributeIDsMutex.RUnlock()
+	return len(fake.getAllAsyncAttributeIDsArgsForCall)
+}
+
+func (fake *FakeLocalParticipant) GetAllAsyncAttributeIDsCalls(stub func() []*livekit.DataTrackSchemaId) {
+	fake.getAllAsyncAttributeIDsMutex.Lock()
+	defer fake.getAllAsyncAttributeIDsMutex.Unlock()
+	fake.GetAllAsyncAttributeIDsStub = stub
+}
+
+func (fake *FakeLocalParticipant) GetAllAsyncAttributeIDsReturns(result1 []*livekit.DataTrackSchemaId) {
+	fake.getAllAsyncAttributeIDsMutex.Lock()
+	defer fake.getAllAsyncAttributeIDsMutex.Unlock()
+	fake.GetAllAsyncAttributeIDsStub = nil
+	fake.getAllAsyncAttributeIDsReturns = struct {
+		result1 []*livekit.DataTrackSchemaId
+	}{result1}
+}
+
+func (fake *FakeLocalParticipant) GetAllAsyncAttributeIDsReturnsOnCall(i int, result1 []*livekit.DataTrackSchemaId) {
+	fake.getAllAsyncAttributeIDsMutex.Lock()
+	defer fake.getAllAsyncAttributeIDsMutex.Unlock()
+	fake.GetAllAsyncAttributeIDsStub = nil
+	if fake.getAllAsyncAttributeIDsReturnsOnCall == nil {
+		fake.getAllAsyncAttributeIDsReturnsOnCall = make(map[int]struct {
+			result1 []*livekit.DataTrackSchemaId
+		})
+	}
+	fake.getAllAsyncAttributeIDsReturnsOnCall[i] = struct {
+		result1 []*livekit.DataTrackSchemaId
+	}{result1}
+}
+
 func (fake *FakeLocalParticipant) GetAnswer() (webrtc.SessionDescription, uint32, error) {
 	fake.getAnswerMutex.Lock()
 	ret, specificReturn := fake.getAnswerReturnsOnCall[len(fake.getAnswerArgsForCall)]