Relay timeouts (Layr-Labs#918)

Signed-off-by: Cody Littley <[email protected]>
cody-littley · Nov 25, 2024 · ac7ffdd · ac7ffdd
1 parent 0a4e852
commit ac7ffdd
Show file tree

Hide file tree

Showing 18 changed files with 547 additions and 120 deletions.
diff --git a/common/aws/cli.go b/common/aws/cli.go
@@ -37,12 +37,6 @@ type ClientConfig struct {
 	// FragmentParallelismConstant helps determine the size of the pool of workers to help upload/download files.
 	// A non-zero value for this parameter adds a constant number of workers. Default is 0.
 	FragmentParallelismConstant int
-	// FragmentReadTimeout is used to bound the maximum time to wait for a single fragmented read.
-	// Default is 30 seconds.
-	FragmentReadTimeout time.Duration
-	// FragmentWriteTimeout is used to bound the maximum time to wait for a single fragmented write.
-	// Default is 30 seconds.
-	FragmentWriteTimeout time.Duration
 }
 
 func ClientFlags(envPrefix string, flagPrefix string) []cli.Flag {
@@ -120,8 +114,6 @@ func ReadClientConfig(ctx *cli.Context, flagPrefix string) ClientConfig {
 		EndpointURL:                 ctx.GlobalString(common.PrefixFlag(flagPrefix, EndpointURLFlagName)),
 		FragmentParallelismFactor:   ctx.GlobalInt(common.PrefixFlag(flagPrefix, FragmentParallelismFactorFlagName)),
 		FragmentParallelismConstant: ctx.GlobalInt(common.PrefixFlag(flagPrefix, FragmentParallelismConstantFlagName)),
-		FragmentReadTimeout:         ctx.GlobalDuration(common.PrefixFlag(flagPrefix, FragmentReadTimeoutFlagName)),
-		FragmentWriteTimeout:        ctx.GlobalDuration(common.PrefixFlag(flagPrefix, FragmentWriteTimeoutFlagName)),
 	}
 }
 
@@ -131,7 +123,5 @@ func DefaultClientConfig() *ClientConfig {
 		Region:                      "us-east-2",
 		FragmentParallelismFactor:   8,
 		FragmentParallelismConstant: 0,
-		FragmentReadTimeout:         30 * time.Second,
-		FragmentWriteTimeout:        30 * time.Second,
 	}
 }
diff --git a/common/aws/s3/client.go b/common/aws/s3/client.go
@@ -241,9 +241,6 @@ func (s *client) FragmentedUploadObject(
 	}
 	resultChannel := make(chan error, len(fragments))
 
-	ctx, cancel := context.WithTimeout(ctx, s.cfg.FragmentWriteTimeout)
-	defer cancel()
-
 	for _, fragment := range fragments {
 		fragmentCapture := fragment
 		s.concurrencyLimiter <- struct{}{}
@@ -301,9 +298,6 @@ func (s *client) FragmentedDownloadObject(
 	}
 	resultChannel := make(chan *readResult, len(fragmentKeys))
 
-	ctx, cancel := context.WithTimeout(ctx, s.cfg.FragmentWriteTimeout)
-	defer cancel()
-
 	for i, fragmentKey := range fragmentKeys {
 		boundFragmentKey := fragmentKey
 		boundI := i

diff --git a/relay/auth/authenticator.go b/relay/auth/authenticator.go
@@ -17,6 +17,7 @@ type RequestAuthenticator interface {
 	// The origin is the address of the peer that sent the request. This may be used to cache auth results
 	// in order to save server resources.
 	AuthenticateGetChunksRequest(
+		ctx context.Context,
 		origin string,
 		request *pb.GetChunksRequest,
 		now time.Time) error
@@ -53,6 +54,7 @@ type requestAuthenticator struct {
 
 // NewRequestAuthenticator creates a new RequestAuthenticator.
 func NewRequestAuthenticator(
+	ctx context.Context,
 	ics core.IndexedChainState,
 	keyCacheSize int,
 	authenticationTimeoutDuration time.Duration) (RequestAuthenticator, error) {
@@ -70,20 +72,20 @@ func NewRequestAuthenticator(
 		keyCache:                      keyCache,
 	}
 
-	err = authenticator.preloadCache()
+	err = authenticator.preloadCache(ctx)
 	if err != nil {
 		return nil, fmt.Errorf("failed to preload cache: %w", err)
 	}
 
 	return authenticator, nil
 }
 
-func (a *requestAuthenticator) preloadCache() error {
+func (a *requestAuthenticator) preloadCache(ctx context.Context) error {
 	blockNumber, err := a.ics.GetCurrentBlockNumber()
 	if err != nil {
 		return fmt.Errorf("failed to get current block number: %w", err)
 	}
-	operators, err := a.ics.GetIndexedOperators(context.Background(), blockNumber)
+	operators, err := a.ics.GetIndexedOperators(ctx, blockNumber)
 	if err != nil {
 		return fmt.Errorf("failed to get operators: %w", err)
 	}
@@ -96,6 +98,7 @@ func (a *requestAuthenticator) preloadCache() error {
 }
 
 func (a *requestAuthenticator) AuthenticateGetChunksRequest(
+	ctx context.Context,
 	origin string,
 	request *pb.GetChunksRequest,
 	now time.Time) error {
@@ -105,7 +108,7 @@ func (a *requestAuthenticator) AuthenticateGetChunksRequest(
 		return nil
 	}
 
-	key, err := a.getOperatorKey(core.OperatorID(request.OperatorId))
+	key, err := a.getOperatorKey(ctx, core.OperatorID(request.OperatorId))
 	if err != nil {
 		return fmt.Errorf("failed to get operator key: %w", err)
 	}
@@ -131,7 +134,7 @@ func (a *requestAuthenticator) AuthenticateGetChunksRequest(
 }
 
 // getOperatorKey returns the public key of the operator with the given ID, caching the result.
-func (a *requestAuthenticator) getOperatorKey(operatorID core.OperatorID) (*core.G2Point, error) {
+func (a *requestAuthenticator) getOperatorKey(ctx context.Context, operatorID core.OperatorID) (*core.G2Point, error) {
 	key, ok := a.keyCache.Get(operatorID)
 	if ok {
 		return key, nil
@@ -141,7 +144,7 @@ func (a *requestAuthenticator) getOperatorKey(operatorID core.OperatorID) (*core
 	if err != nil {
 		return nil, fmt.Errorf("failed to get current block number: %w", err)
 	}
-	operators, err := a.ics.GetIndexedOperators(context.Background(), blockNumber)
+	operators, err := a.ics.GetIndexedOperators(ctx, blockNumber)
 	if err != nil {
 		return nil, fmt.Errorf("failed to get operators: %w", err)
 	}

diff --git a/relay/auth/authenticator_test.go b/relay/auth/authenticator_test.go
@@ -15,6 +15,8 @@ import (
 func TestMockSigning(t *testing.T) {
 	tu.InitializeRandom()
 
+	ctx := context.Background()
+
 	operatorID := mock.MakeOperatorId(0)
 	stakes := map[core.QuorumID]map[core.OperatorID]int{
 		core.QuorumID(0): {
@@ -24,7 +26,7 @@ func TestMockSigning(t *testing.T) {
 	ics, err := mock.NewChainDataMock(stakes)
 	require.NoError(t, err)
 
-	operators, err := ics.GetIndexedOperators(context.Background(), 0)
+	operators, err := ics.GetIndexedOperators(ctx, 0)
 	require.NoError(t, err)
 
 	operator, ok := operators[operatorID]
@@ -46,6 +48,8 @@ func TestMockSigning(t *testing.T) {
 func TestValidRequest(t *testing.T) {
 	tu.InitializeRandom()
 
+	ctx := context.Background()
+
 	operatorID := mock.MakeOperatorId(0)
 	stakes := map[core.QuorumID]map[core.OperatorID]int{
 		core.QuorumID(0): {
@@ -58,7 +62,7 @@ func TestValidRequest(t *testing.T) {
 
 	timeout := 10 * time.Second
 
-	authenticator, err := NewRequestAuthenticator(ics, 1024, timeout)
+	authenticator, err := NewRequestAuthenticator(ctx, ics, 1024, timeout)
 	require.NoError(t, err)
 
 	request := randomGetChunksRequest()
@@ -69,6 +73,7 @@ func TestValidRequest(t *testing.T) {
 	now := time.Now()
 
 	err = authenticator.AuthenticateGetChunksRequest(
+		ctx,
 		"foobar",
 		request,
 		now)
@@ -83,12 +88,14 @@ func TestValidRequest(t *testing.T) {
 	start := now
 	for now.Before(start.Add(timeout)) {
 		err = authenticator.AuthenticateGetChunksRequest(
+			ctx,
 			"foobar",
 			invalidRequest,
 			now)
 		require.NoError(t, err)
 
 		err = authenticator.AuthenticateGetChunksRequest(
+			ctx,
 			"baz",
 			invalidRequest,
 			now)
@@ -99,6 +106,7 @@ func TestValidRequest(t *testing.T) {
 
 	// After the timeout elapses, new requests should trigger authentication.
 	err = authenticator.AuthenticateGetChunksRequest(
+		ctx,
 		"foobar",
 		invalidRequest,
 		now)
@@ -108,6 +116,8 @@ func TestValidRequest(t *testing.T) {
 func TestAuthenticationSavingDisabled(t *testing.T) {
 	tu.InitializeRandom()
 
+	ctx := context.Background()
+
 	operatorID := mock.MakeOperatorId(0)
 	stakes := map[core.QuorumID]map[core.OperatorID]int{
 		core.QuorumID(0): {
@@ -121,7 +131,7 @@ func TestAuthenticationSavingDisabled(t *testing.T) {
 	// This disables saving of authentication results.
 	timeout := time.Duration(0)
 
-	authenticator, err := NewRequestAuthenticator(ics, 1024, timeout)
+	authenticator, err := NewRequestAuthenticator(ctx, ics, 1024, timeout)
 	require.NoError(t, err)
 
 	request := randomGetChunksRequest()
@@ -132,6 +142,7 @@ func TestAuthenticationSavingDisabled(t *testing.T) {
 	now := time.Now()
 
 	err = authenticator.AuthenticateGetChunksRequest(
+		ctx,
 		"foobar",
 		request,
 		now)
@@ -144,6 +155,7 @@ func TestAuthenticationSavingDisabled(t *testing.T) {
 	invalidRequest.OperatorSignature = signature // the previous signature is invalid here
 
 	err = authenticator.AuthenticateGetChunksRequest(
+		ctx,
 		"foobar",
 		invalidRequest,
 		now)
@@ -153,6 +165,8 @@ func TestAuthenticationSavingDisabled(t *testing.T) {
 func TestNonExistingClient(t *testing.T) {
 	tu.InitializeRandom()
 
+	ctx := context.Background()
+
 	operatorID := mock.MakeOperatorId(0)
 	stakes := map[core.QuorumID]map[core.OperatorID]int{
 		core.QuorumID(0): {
@@ -165,7 +179,7 @@ func TestNonExistingClient(t *testing.T) {
 
 	timeout := 10 * time.Second
 
-	authenticator, err := NewRequestAuthenticator(ics, 1024, timeout)
+	authenticator, err := NewRequestAuthenticator(ctx, ics, 1024, timeout)
 	require.NoError(t, err)
 
 	invalidOperatorID := tu.RandomBytes(32)
@@ -174,6 +188,7 @@ func TestNonExistingClient(t *testing.T) {
 	request.OperatorId = invalidOperatorID
 
 	err = authenticator.AuthenticateGetChunksRequest(
+		ctx,
 		"foobar",
 		request,
 		time.Now())
@@ -183,6 +198,8 @@ func TestNonExistingClient(t *testing.T) {
 func TestBadSignature(t *testing.T) {
 	tu.InitializeRandom()
 
+	ctx := context.Background()
+
 	operatorID := mock.MakeOperatorId(0)
 	stakes := map[core.QuorumID]map[core.OperatorID]int{
 		core.QuorumID(0): {
@@ -195,7 +212,7 @@ func TestBadSignature(t *testing.T) {
 
 	timeout := 10 * time.Second
 
-	authenticator, err := NewRequestAuthenticator(ics, 1024, timeout)
+	authenticator, err := NewRequestAuthenticator(ctx, ics, 1024, timeout)
 	require.NoError(t, err)
 
 	request := randomGetChunksRequest()
@@ -205,6 +222,7 @@ func TestBadSignature(t *testing.T) {
 	now := time.Now()
 
 	err = authenticator.AuthenticateGetChunksRequest(
+		ctx,
 		"foobar",
 		request,
 		now)
@@ -217,6 +235,7 @@ func TestBadSignature(t *testing.T) {
 	request.OperatorSignature[0] = request.OperatorSignature[0] ^ 1
 
 	err = authenticator.AuthenticateGetChunksRequest(
+		ctx,
 		"foobar",
 		request,
 		now)

diff --git a/relay/blob_provider.go b/relay/blob_provider.go
@@ -7,6 +7,7 @@ import (
 	"github.com/Layr-Labs/eigenda/disperser/common/v2/blobstore"
 	"github.com/Layr-Labs/eigenda/relay/cache"
 	"github.com/Layr-Labs/eigensdk-go/logging"
+	"time"
 )
 
 // blobProvider encapsulates logic for fetching blobs. Utilized by the relay Server.
@@ -20,6 +21,9 @@ type blobProvider struct {
 
 	// blobCache is an LRU cache of blobs.
 	blobCache cache.CachedAccessor[v2.BlobKey, []byte]
+
+	// fetchTimeout is the maximum time to wait for a blob fetch operation to complete.
+	fetchTimeout time.Duration
 }
 
 // newBlobProvider creates a new blobProvider.
@@ -28,12 +32,14 @@ func newBlobProvider(
 	logger logging.Logger,
 	blobStore *blobstore.BlobStore,
 	blobCacheSize int,
-	maxIOConcurrency int) (*blobProvider, error) {
+	maxIOConcurrency int,
+	fetchTimeout time.Duration) (*blobProvider, error) {
 
 	server := &blobProvider{
-		ctx:       ctx,
-		logger:    logger,
-		blobStore: blobStore,
+		ctx:          ctx,
+		logger:       logger,
+		blobStore:    blobStore,
+		fetchTimeout: fetchTimeout,
 	}
 
 	c, err := cache.NewCachedAccessor[v2.BlobKey, []byte](blobCacheSize, maxIOConcurrency, server.fetchBlob)
@@ -46,9 +52,8 @@ func newBlobProvider(
 }
 
 // GetBlob retrieves a blob from the blob store.
-func (s *blobProvider) GetBlob(blobKey v2.BlobKey) ([]byte, error) {
-
-	data, err := s.blobCache.Get(blobKey)
+func (s *blobProvider) GetBlob(ctx context.Context, blobKey v2.BlobKey) ([]byte, error) {
+	data, err := s.blobCache.Get(ctx, blobKey)
 
 	if err != nil {
 		// It should not be possible for external users to force an error here since we won't
@@ -62,7 +67,10 @@ func (s *blobProvider) GetBlob(blobKey v2.BlobKey) ([]byte, error) {
 
 // fetchBlob retrieves a single blob from the blob store.
 func (s *blobProvider) fetchBlob(blobKey v2.BlobKey) ([]byte, error) {
-	data, err := s.blobStore.GetBlob(s.ctx, blobKey)
+	ctx, cancel := context.WithTimeout(s.ctx, s.fetchTimeout)
+	defer cancel()
+
+	data, err := s.blobStore.GetBlob(ctx, blobKey)
 	if err != nil {
 		s.logger.Errorf("Failed to fetch blob: %v", err)
 		return nil, err