Added blob table data structure.

tools/traffic/table/blob_metadata.go

@@ -0,0 +1,71 @@
+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
+
+	// index describes the position of this blob within the blobTable.
+	index uint
+}
+
+// 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,
+		index:                0,
+	}
+}
+
+// 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_table.go

@@ -0,0 +1,140 @@
+package table
+
+import (
+	"fmt"
+	"math/rand"
+	"sync"
+)
+
+// BlobTable tracks blobs written by the traffic generator. This is a thread safe data structure.
+type BlobTable struct {
+
+	// blobs contains all blobs currently tracked by the requiredReads.
+	blobs []*BlobMetadata
+
+	// size describes the total number of blobs currently tracked by the requiredReads.
+	// size may be smaller than the capacity of the blobs slice.
+	size uint
+
+	// lock is used to synchronize access to the requiredReads.
+	lock sync.RWMutex
+}
+
+// NewBlobTable creates a new BlobTable instance.
+func NewBlobTable() BlobTable {
+	return BlobTable{
+		blobs: make([]*BlobMetadata, 0),
+		size:  0,
+	}
+}
+
+// Size returns the total number of blobs currently tracked by the requiredReads.
+func (table *BlobTable) Size() uint {
+	table.lock.RLock()
+	defer table.lock.RUnlock()
+
+	return table.size
+}
+
+// Get returns the blob at the specified index. Returns nil if the index is out of bounds.
+func (table *BlobTable) Get(index uint) *BlobMetadata {
+	table.lock.RLock()
+	defer table.lock.RUnlock()
+
+	if index >= table.size {
+		return nil
+	}
+
+	return table.blobs[index]
+}
+
+// Add a blob to the requiredReads.
+func (table *BlobTable) Add(blob *BlobMetadata) {
+	table.lock.Lock()
+	defer table.lock.Unlock()
+
+	blob.index = table.size
+
+	if table.size == uint(len(table.blobs)) {
+		table.blobs = append(table.blobs, blob)
+	} else {
+		table.blobs[table.size] = blob
+	}
+
+	table.size++
+}
+
+// AddOrReplace is equivalent to Add if there is capacity, or replaces an existing blob at random
+// if the is no remaining capacity. This method is a no-op if maximumCapacity is 0.
+func (table *BlobTable) AddOrReplace(blob *BlobMetadata, maximumCapacity uint) {
+	if maximumCapacity == 0 {
+		return
+	}
+
+	table.lock.Lock()
+	defer table.lock.Unlock()
+
+	if table.size >= maximumCapacity {
+		// replace random existing blob
+		index := rand.Int31n(int32(table.size))
+		table.blobs[index] = blob
+		blob.index = uint(index)
+	} else {
+		// add new blob
+		blob.index = table.size
+		if table.size == uint(len(table.blobs)) {
+			table.blobs = append(table.blobs, blob)
+		} else {
+			table.blobs[table.size] = blob
+		}
+		table.size++
+	}
+}
+
+// GetRandom returns a random blob currently tracked by the requiredReads. Returns nil if the requiredReads is empty.
+// Optionally decrements the read  permits of the blob if decrement is true. If the number of read permits
+// reaches 0, the blob is removed  from the requiredReads. Returns the blob metadata (if there is at least one blob
+// in the table) and a boolean indicating whether the blob was removed from the table as a result of this operation.
+func (table *BlobTable) GetRandom(decrement bool) (*BlobMetadata, bool) {
+	if decrement {
+		table.lock.Lock()
+		defer table.lock.Unlock()
+	} else {
+		table.lock.RLock()
+		defer table.lock.RUnlock()
+	}
+
+	if table.size == 0 {
+		return nil, false
+	}
+
+	blob := table.blobs[rand.Int31n(int32(table.size))]
+
+	removed := false
+	if decrement && blob.remainingReadPermits != -1 {
+		blob.remainingReadPermits--
+		if blob.remainingReadPermits == 0 {
+			table.remove(blob)
+			removed = true
+		}
+	}
+
+	return blob, removed
+}
+
+// remove a blob from the requiredReads.
+func (table *BlobTable) remove(blob *BlobMetadata) {
+	if table.blobs[blob.index] != blob {
+		panic(fmt.Sprintf("blob %x is not not present in the requiredReads at index %d", blob.Key(), blob.index))
+	}
+
+	if table.size == 1 {
+		table.blobs[0] = nil
+	} else {
+		// Move the last blob to the position of the blob being removed.
+		table.blobs[blob.index] = table.blobs[table.size-1]
+		table.blobs[blob.index].index = blob.index
+		table.blobs[table.size-1] = nil
+	}
+	table.size--
+}

tools/traffic/table/blob_table_test.go

@@ -0,0 +1,175 @@
+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()
+
+	table := NewBlobTable()
+	assert.Equal(t, uint(0), table.Size())
+
+	size := 1024
+	expectedMetadata := make([]*BlobMetadata, 0)
+	for i := 0; i < size; i++ {
+		metadata := randomMetadata(1)
+		table.Add(metadata)
+		expectedMetadata = append(expectedMetadata, metadata)
+		assert.Equal(t, uint(i+1), table.Size())
+	}
+
+	for i := 0; i < size; i++ {
+		assert.Equal(t, expectedMetadata[i], table.Get(uint(i)))
+	}
+
+	// Requesting an index that is out of bounds should return nil.
+	assert.Nil(t, table.Get(uint(size)))
+}
+
+// TestGetRandomWithRemoval tests getting a random blob data, but where the number of permits per blob is unlimited.
+func TestGetRandomNoRemovalByConfiguration(t *testing.T) {
+	tu.InitializeRandom()
+
+	table := NewBlobTable()
+	assert.Equal(t, uint(0), table.Size())
+
+	// Requesting a random element from an empty table should return nil.
+	element, _ := table.GetRandom(true)
+	assert.Nil(t, element)
+
+	expectedMetadata := make([]*BlobMetadata, 0)
+	size := 128
+	for i := 0; i < size; i++ {
+		metadata := randomMetadata(-1) // -1 == unlimited permits
+		table.Add(metadata)
+		expectedMetadata = append(expectedMetadata, metadata)
+		assert.Equal(t, uint(i+1), table.Size())
+	}
+
+	randomIndices := make(map[uint]bool)
+
+	// Query more times than the number of blobs to ensure that blobs are not removed.
+	for i := 0; i < size*8; i++ {
+		// This parameter will be ignored given that the number of permits is unlimited.
+		// But not a bad thing to exercise the code path.
+		decrement := rand.Intn(2) == 1
+
+		metadata, removed := table.GetRandom(decrement)
+		assert.False(t, removed)
+		assert.NotNil(t, metadata)
+		assert.Equal(t, expectedMetadata[metadata.index], metadata)
+
+		randomIndices[metadata.index] = 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 := NewBlobTable()
+	assert.Equal(t, uint(0), table.Size())
+
+	// Requesting a random element from an empty table should return nil.
+	element, _ := table.GetRandom(true)
+	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 random elements without decrementing should not remove any elements.
+	for i := 0; i < size; i++ {
+		metadata, removed := table.GetRandom(false)
+		assert.NotNil(t, metadata)
+		_, exists := expectedMetadata[string(metadata.Key())]
+		assert.True(t, exists)
+		assert.False(t, removed)
+	}
+	assert.Equal(t, uint(size), 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, removed := table.GetRandom(true)
+		assert.NotNil(t, metadata)
+
+		k := string(metadata.Key())
+		permitsUsed := expectedMetadata[k] + 1
+		expectedMetadata[k] = permitsUsed
+		assert.LessOrEqual(t, permitsUsed, uint(permitCount))
+
+		if int(permitsUsed) == permitCount {
+			assert.True(t, removed)
+		} else {
+			assert.False(t, removed)
+		}
+	}
+
+	assert.Equal(t, uint(0), table.Size())
+}
+
+// TestAddOrReplace tests adding blobs to a table with a maximum capacity. The table should replace blobs when full.
+func TestAddOrReplace(t *testing.T) {
+	tu.InitializeRandom()
+
+	table := NewBlobTable()
+	assert.Equal(t, uint(0), table.Size())
+
+	// Adding data to a table with capacity 0 should be a no-op.
+	table.AddOrReplace(randomMetadata(1), 0)
+	assert.Equal(t, uint(0), table.Size())
+
+	randomIndices := make(map[uint]bool)
+
+	size := 1024
+	for i := 0; i < size*2; i++ {
+		metadata := randomMetadata(-1) // -1 == unlimited permits
+
+		initialSize := table.Size()
+		table.AddOrReplace(metadata, uint(size))
+		resultingSize := table.Size()
+
+		assert.LessOrEqual(t, resultingSize, uint(size))
+		if initialSize < uint(size) {
+			assert.Equal(t, initialSize+1, resultingSize)
+		} else {
+			randomIndices[metadata.index] = true
+		}
+
+		// Verify that the metadata is in the table.
+		assert.Less(t, metadata.index, table.Size())
+		assert.Equal(t, metadata, table.Get(metadata.index))
+	}
+
+	// Sanity check: ensure that replacements happened at least 10 different indices. This check is attempting to
+	// verify that we are actually replacing blobs. The probability of this check failing is extremely low if
+	// the random number generator is working correctly.
+	assert.GreaterOrEqual(t, len(randomIndices), 10)
+}