Added blob table data structure.

tools/traffic/table/blob_metadata.go

@@ -0,0 +1,67 @@
+package table
+
+// BlobMetadata encapsulates various information about a blob written by the traffic generator.
+type BlobMetadata struct {
+	// key of the blob, set when the blob is initially uploaded.
+	key []byte
+
+	// checksum of the blob.
+	checksum *[16]byte
+
+	// size of the blob, in bytes.
+	size uint
+
+	// blobIndex of the blob.
+	blobIndex uint
+
+	// remainingReadPermits describes the maximum number of remaining reads permitted against this blob.
+	// If -1 then an unlimited number of reads are permitted.
+	remainingReadPermits int
+}
+
+// NewBlobMetadata creates a new BlobMetadata instance. The readPermits parameter describes the maximum number of
+// remaining reads permitted against this blob. If -1 then an unlimited number of reads are permitted.
+func NewBlobMetadata(
+	key []byte,
+	checksum *[16]byte,
+	size uint,
+	blobIndex uint,
+	readPermits int) *BlobMetadata {
+
+	if readPermits == 0 {
+		panic("readPermits must be greater than 0, or -1 for unlimited reads")
+	}
+
+	return &BlobMetadata{
+		key:                  key,
+		checksum:             checksum,
+		size:                 size,
+		blobIndex:            blobIndex,
+		remainingReadPermits: readPermits,
+	}
+}
+
+// Key returns the key of the blob.
+func (blob *BlobMetadata) Key() []byte {
+	return blob.key
+}
+
+// Checksum returns the checksum of the blob.
+func (blob *BlobMetadata) Checksum() *[16]byte {
+	return blob.checksum
+}
+
+// Size returns the size of the blob, in bytes.
+func (blob *BlobMetadata) Size() uint {
+	return blob.size
+}
+
+// BlobIndex returns the blobIndex of the blob.
+func (blob *BlobMetadata) BlobIndex() uint {
+	return blob.blobIndex
+}
+
+// RemainingReadPermits returns the maximum number of remaining reads permitted against this blob.
+func (blob *BlobMetadata) RemainingReadPermits() int {
+	return blob.remainingReadPermits
+}

tools/traffic/table/blob_store.go

@@ -0,0 +1,62 @@
+package table
+
+import "sync"
+
+// BlobStore is a thread safe data structure that tracks blobs written by the traffic generator.
+type BlobStore struct {
+
+	// blobs contains all blobs currently tracked by the store.
+	blobs map[uint64]*BlobMetadata
+
+	// nextKey describes the next key to used for the blobs map.
+	nextKey uint64
+
+	lock sync.Mutex
+}
+
+// NewBlobStore creates a new BlobStore instance.
+func NewBlobStore() *BlobStore {
+	return &BlobStore{
+		blobs:   make(map[uint64]*BlobMetadata),
+		nextKey: 0,
+	}
+}
+
+// Add a blob to the store.
+func (store *BlobStore) Add(blob *BlobMetadata) {
+	store.lock.Lock()
+	defer store.lock.Unlock()
+
+	store.blobs[store.nextKey] = blob
+	store.nextKey++
+}
+
+// GetNext returns the next blob in the store. Decrements the blob's read permits, removing it
+// from the store if the permits are exhausted. This method makes no guarantees about the order
+// in which blobs are returned. Returns nil if no blobs are available.
+func (store *BlobStore) GetNext() *BlobMetadata {
+	store.lock.Lock()
+	defer store.lock.Unlock()
+
+	for key, blob := range store.blobs {
+		// Always return the first blob found.
+
+		if blob.remainingReadPermits > 0 {
+			blob.remainingReadPermits--
+			if blob.remainingReadPermits == 0 {
+				delete(store.blobs, key)
+			}
+		}
+
+		return blob
+	}
+	return nil
+}
+
+// Size returns the number of blobs currently stored.
+func (store *BlobStore) Size() uint {
+	store.lock.Lock()
+	defer store.lock.Unlock()
+
+	return uint(len(store.blobs))
+}

tools/traffic/table/blob_store_test.go

@@ -0,0 +1,111 @@
+package table
+
+import (
+	tu "github.com/Layr-Labs/eigenda/common/testutils"
+	"github.com/stretchr/testify/assert"
+	"golang.org/x/exp/rand"
+	"testing"
+)
+
+// randomMetadata generates a random BlobMetadata instance.
+func randomMetadata(permits int) *BlobMetadata {
+	key := make([]byte, 32)
+	checksum := [16]byte{}
+	_, _ = rand.Read(key)
+	_, _ = rand.Read(checksum[:])
+	return NewBlobMetadata(key, &checksum, 1024, 0, permits)
+}
+
+// TestBasicOperation tests basic operations of the BlobTable. Adds blobs and iterates over them.
+func TestBasicOperation(t *testing.T) {
+	tu.InitializeRandom()
+
+	store := NewBlobStore()
+	assert.Equal(t, uint(0), store.Size())
+
+	size := 1024
+	expectedMetadata := make([]*BlobMetadata, 0)
+	for i := 0; i < size; i++ {
+		metadata := randomMetadata(1)
+		store.Add(metadata)
+		expectedMetadata = append(expectedMetadata, metadata)
+		assert.Equal(t, uint(i+1), store.Size())
+	}
+
+	for i := 0; i < size; i++ {
+		assert.Equal(t, expectedMetadata[i], store.blobs[uint64(i)])
+	}
+}
+
+// TestGetRandomWithRemoval tests getting a random blob data, but where the number of permits per blob is unlimited.
+func TestGetRandomNoRemoval(t *testing.T) {
+	tu.InitializeRandom()
+
+	table := NewBlobStore()
+	assert.Equal(t, uint(0), table.Size())
+
+	// Requesting a random element from an empty table should return nil.
+	element := table.GetNext()
+	assert.Nil(t, element)
+
+	expectedMetadata := make(map[string]*BlobMetadata)
+	size := 128
+	for i := 0; i < size; i++ {
+		metadata := randomMetadata(-1) // -1 == unlimited permits
+		table.Add(metadata)
+		expectedMetadata[string(metadata.key)] = metadata
+		assert.Equal(t, uint(i+1), table.Size())
+	}
+
+	randomIndices := make(map[string]bool)
+
+	// Query more times than the number of blobs to ensure that blobs are not removed.
+	for i := 0; i < size*8; i++ {
+		metadata := table.GetNext()
+		assert.NotNil(t, metadata)
+		assert.Equal(t, expectedMetadata[string(metadata.key)], metadata)
+		randomIndices[string(metadata.key)] = true
+	}
+
+	// Sanity check: ensure that at least 10 different blobs were returned. This check is attempting to verify
+	// that we are actually getting random blobs. The probability of this check failing is extremely low if
+	// the random number generator is working correctly.
+	assert.GreaterOrEqual(t, len(randomIndices), 10)
+}
+
+// TestGetRandomWithRemoval tests getting a random blob data, where the number of permits per blob is limited.
+func TestGetRandomWithRemoval(t *testing.T) {
+	tu.InitializeRandom()
+
+	table := NewBlobStore()
+	assert.Equal(t, uint(0), table.Size())
+
+	// Requesting a random element from an empty table should return nil.
+	element := table.GetNext()
+	assert.Nil(t, element)
+
+	permitCount := 2
+
+	size := 1024
+	expectedMetadata := make(map[string]uint)
+	for i := 0; i < size; i++ {
+		metadata := randomMetadata(permitCount)
+		table.Add(metadata)
+		expectedMetadata[string(metadata.Key())] = 0
+		assert.Equal(t, uint(i+1), table.Size())
+	}
+
+	// Requesting elements a number of times equal to the size times the number of permits should completely
+	// drain the table and return all elements a number of times equal to the number of permits.
+	for i := 0; i < size*permitCount; i++ {
+		metadata := table.GetNext()
+		assert.NotNil(t, metadata)
+
+		k := string(metadata.Key())
+		permitsUsed := expectedMetadata[k] + 1
+		expectedMetadata[k] = permitsUsed
+		assert.LessOrEqual(t, permitsUsed, uint(permitCount))
+	}
+
+	assert.Equal(t, uint(0), table.Size())
+}