storage: source storage timeout

oasisprotocol · Jan 7, 2023 · 993da0b · 993da0b
1 parent b13917e
commit 993da0b
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 0 deletions.
diff --git a/storage/oasis/consensus.go b/storage/oasis/consensus.go
@@ -3,6 +3,7 @@ package oasis
 import (
 	"context"
 	"fmt"
+	"time"
 
 	"github.com/oasisprotocol/oasis-core/go/beacon/api"
 	"github.com/oasisprotocol/oasis-core/go/common"
@@ -62,6 +63,9 @@ func (cc *ConsensusClient) Name() string {
 
 // GetEpoch returns the epoch number at the specified block height.
 func (cc *ConsensusClient) GetEpoch(ctx context.Context, height int64) (api.EpochTime, error) {
+	ctx, cancel := context.WithTimeout(ctx, 61*time.Second)
+	defer cancel()
+
 	return cc.client.Beacon().GetEpoch(ctx, height)
 }
 
@@ -110,6 +114,9 @@ func (cc *ConsensusClient) AllData(ctx context.Context, height int64) (*storage.
 
 // BlockData retrieves data about a consensus block at the provided block height.
 func (cc *ConsensusClient) BlockData(ctx context.Context, height int64) (*storage.ConsensusBlockData, error) {
+	ctx, cancel := context.WithTimeout(ctx, 61*time.Second)
+	defer cancel()
+
 	block, err := cc.client.GetBlock(ctx, height)
 	if err != nil {
 		return nil, err
@@ -145,6 +152,9 @@ func (cc *ConsensusClient) BlockData(ctx context.Context, height int64) (*storag
 
 // BeaconData retrieves the beacon for the provided block height.
 func (cc *ConsensusClient) BeaconData(ctx context.Context, height int64) (*storage.BeaconData, error) {
+	ctx, cancel := context.WithTimeout(ctx, 61*time.Second)
+	defer cancel()
+
 	// NOTE: The random beacon endpoint is in flux.
 	// GetBeacon() consistently errors out (at least for heights soon after Damask genesis) with "beacon: random beacon not available".
 	// beacon, err := cc.client.Beacon().GetBeacon(ctx, height)
@@ -166,6 +176,9 @@ func (cc *ConsensusClient) BeaconData(ctx context.Context, height int64) (*stora
 
 // RegistryData retrieves registry events at the provided block height.
 func (cc *ConsensusClient) RegistryData(ctx context.Context, height int64) (*storage.RegistryData, error) {
+	ctx, cancel := context.WithTimeout(ctx, 61*time.Second)
+	defer cancel()
+
 	events, err := cc.client.Registry().GetEvents(ctx, height)
 	if err != nil {
 		return nil, err
@@ -201,6 +214,9 @@ func (cc *ConsensusClient) RegistryData(ctx context.Context, height int64) (*sto
 
 // StakingData retrieves staking events at the provided block height.
 func (cc *ConsensusClient) StakingData(ctx context.Context, height int64) (*storage.StakingData, error) {
+	ctx, cancel := context.WithTimeout(ctx, 61*time.Second)
+	defer cancel()
+
 	events, err := cc.client.Staking().GetEvents(ctx, height)
 	if err != nil {
 		return nil, err
@@ -242,6 +258,9 @@ func (cc *ConsensusClient) StakingData(ctx context.Context, height int64) (*stor
 
 // SchedulerData retrieves validators and runtime committees at the provided block height.
 func (cc *ConsensusClient) SchedulerData(ctx context.Context, height int64) (*storage.SchedulerData, error) {
+	ctx, cancel := context.WithTimeout(ctx, 61*time.Second)
+	defer cancel()
+
 	validators, err := cc.client.Scheduler().GetValidators(ctx, height)
 	if err != nil {
 		return nil, err
@@ -274,6 +293,9 @@ func (cc *ConsensusClient) SchedulerData(ctx context.Context, height int64) (*st
 
 // GovernanceData retrieves governance events at the provided block height.
 func (cc *ConsensusClient) GovernanceData(ctx context.Context, height int64) (*storage.GovernanceData, error) {
+	ctx, cancel := context.WithTimeout(ctx, 61*time.Second)
+	defer cancel()
+
 	events, err := cc.client.Governance().GetEvents(ctx, height)
 	if err != nil {
 		return nil, err
@@ -322,6 +344,9 @@ func (cc *ConsensusClient) GovernanceData(ctx context.Context, height int64) (*s
 
 // RootHashData retrieves roothash events at the provided block height.
 func (cc *ConsensusClient) RootHashData(ctx context.Context, height int64) (*storage.RootHashData, error) {
+	ctx, cancel := context.WithTimeout(ctx, 61*time.Second)
+	defer cancel()
+
 	events, err := cc.client.RootHash().GetEvents(ctx, height)
 	if err != nil {
 		return nil, err

diff --git a/storage/oasis/runtime.go b/storage/oasis/runtime.go
@@ -3,6 +3,7 @@ package oasis
 import (
 	"context"
 	"fmt"
+	"time"
 
 	config "github.com/oasisprotocol/oasis-sdk/client-sdk/go/config"
 	connection "github.com/oasisprotocol/oasis-sdk/client-sdk/go/connection"
@@ -27,6 +28,9 @@ type RuntimeClient struct {
 
 // BlockData gets block data in the specified round.
 func (rc *RuntimeClient) BlockData(ctx context.Context, round uint64) (*storage.RuntimeBlockData, error) {
+	ctx, cancel := context.WithTimeout(ctx, 61*time.Second)
+	defer cancel()
+
 	block, err := rc.client.GetBlock(ctx, round)
 	if err != nil {
 		return nil, err
@@ -46,6 +50,9 @@ func (rc *RuntimeClient) BlockData(ctx context.Context, round uint64) (*storage.
 
 // CoreData gets data in the specified round emitted by the `core` module.
 func (rc *RuntimeClient) CoreData(ctx context.Context, round uint64) (*storage.CoreData, error) {
+	ctx, cancel := context.WithTimeout(ctx, 61*time.Second)
+	defer cancel()
+
 	events, err := rc.client.Core.GetEvents(ctx, round)
 	if err != nil {
 		return nil, err
@@ -69,6 +76,9 @@ func (rc *RuntimeClient) CoreData(ctx context.Context, round uint64) (*storage.C
 
 // AccountsData gets data in the specified round emitted by the `accounts` module.
 func (rc *RuntimeClient) AccountsData(ctx context.Context, round uint64) (*storage.AccountsData, error) {
+	ctx, cancel := context.WithTimeout(ctx, 61*time.Second)
+	defer cancel()
+
 	events, err := rc.client.Accounts.GetEvents(ctx, round)
 	if err != nil {
 		return nil, err
@@ -100,6 +110,9 @@ func (rc *RuntimeClient) AccountsData(ctx context.Context, round uint64) (*stora
 
 // ConsensusAccountsData gets data in the specified round emitted by the `consensusaccounts` module.
 func (rc *RuntimeClient) ConsensusAccountsData(ctx context.Context, round uint64) (*storage.ConsensusAccountsData, error) {
+	ctx, cancel := context.WithTimeout(ctx, 61*time.Second)
+	defer cancel()
+
 	events, err := rc.client.ConsensusAccounts.GetEvents(ctx, round)
 	if err != nil {
 		return nil, err
@@ -126,6 +139,9 @@ func (rc *RuntimeClient) ConsensusAccountsData(ctx context.Context, round uint64
 }
 
 func (rc *RuntimeClient) EVMSimulateCall(ctx context.Context, round uint64, gasPrice []byte, gasLimit uint64, caller []byte, address []byte, value []byte, data []byte) ([]byte, error) {
+	ctx, cancel := context.WithTimeout(ctx, 61*time.Second)
+	defer cancel()
+
 	return evm.NewV1(rc.client).SimulateCall(ctx, round, gasPrice, gasLimit, caller, address, value, data)
 }