diff --git a/go/oasis-node/cmd/debug/txsource/workload/queries.go b/go/oasis-node/cmd/debug/txsource/workload/queries.go index 31c9e05a3ae..88c3238b562 100644 --- a/go/oasis-node/cmd/debug/txsource/workload/queries.go +++ b/go/oasis-node/cmd/debug/txsource/workload/queries.go @@ -39,6 +39,11 @@ const ( queriesEarliestHeightRatio = 0.1 // Ratio of queries that should query latest available height. queriesLatestHeightRatio = 0.1 + + // queriesIterationTimeout is the combined timeout for running all the queries that are executed + // in a single iteration. The purpose of this timeout is to prevent the client being stuck and + // treating that as an error instead. + queriesIterationTimeout = 60 * time.Second ) // QueriesFlags are the queries workload flags. @@ -422,6 +427,61 @@ func (q *queries) doRuntimeQueries(ctx context.Context, rng *rand.Rand) error { return nil } +func (q *queries) doQueries(ctx context.Context, rng *rand.Rand) error { + block, err := q.consensus.GetBlock(ctx, consensus.HeightLatest) + if err != nil { + return fmt.Errorf("consensus.GetBlock error: %w", err) + } + + // Determine the earliest height that we can query. + earliestHeight := int64(1) + if numKept := viper.GetInt64(CfgConsensusNumKeptVersions); numKept < block.Height { + earliestHeight = block.Height - numKept + } + + // Select height at which queries should be done. Earliest and latest + // heights are special cased with increased probability to be selected. + var height int64 + p := rng.Float32() + switch { + case p < queriesEarliestHeightRatio: + height = earliestHeight + case p < queriesEarliestHeightRatio+queriesLatestHeightRatio: + height = block.Height + default: + // [earliestHeight, block.Height] + height = rng.Int63n(block.Height-earliestHeight+1) + earliestHeight + } + + q.logger.Debug("Doing queries", + "height", height, + "height_latest", block.Height, + ) + + if err := q.doConsensusQueries(ctx, rng, height); err != nil { + return fmt.Errorf("consensus queries error: %w", err) + } + if err := q.doSchedulerQueries(ctx, rng, height); err != nil { + return fmt.Errorf("scheduler queries error: %w", err) + } + if err := q.doRegistryQueries(ctx, rng, height); err != nil { + return fmt.Errorf("registry queries error: %w", err) + } + if err := q.doStakingQueries(ctx, rng, height); err != nil { + return fmt.Errorf("staking queries error: %w", err) + } + if err := q.doRuntimeQueries(ctx, rng); err != nil { + return fmt.Errorf("runtime queries error: %w", err) + } + + q.logger.Debug("Queries done", + "height", height, + "height_latest", block.Height, + ) + + return nil +} + func (q *queries) Run(gracefulExit context.Context, rng *rand.Rand, conn *grpc.ClientConn, cnsc consensus.ClientBackend, fundingAccount signature.Signer) error { ctx := context.Background() @@ -458,57 +518,14 @@ func (q *queries) Run(gracefulExit context.Context, rng *rand.Rand, conn *grpc.C q.runtimeGenesisRound = resp.Header.Round for { - block, err := q.consensus.GetBlock(ctx, consensus.HeightLatest) - if err != nil { - return fmt.Errorf("consensus.GetBlock error: %w", err) - } - - // Determine the earliest height that we can query. - earliestHeight := int64(1) - if numKept := viper.GetInt64(CfgConsensusNumKeptVersions); numKept < block.Height { - earliestHeight = block.Height - numKept - } - - // Select height at which queries should be done. Earliest and latest - // heights are special cased with increased probability to be selected. - var height int64 - p := rng.Float32() - switch { - case p < queriesEarliestHeightRatio: - height = earliestHeight - case p < queriesEarliestHeightRatio+queriesLatestHeightRatio: - height = block.Height - default: - // [earliestHeight, block.Height] - height = rng.Int63n(block.Height-earliestHeight+1) + earliestHeight - } + loopCtx, cancel := context.WithTimeout(ctx, queriesIterationTimeout) - q.logger.Debug("Doing queries", - "height", height, - "height_latest", block.Height, - ) - - if err := q.doConsensusQueries(ctx, rng, height); err != nil { - return fmt.Errorf("consensus queries error: %w", err) - } - if err := q.doSchedulerQueries(ctx, rng, height); err != nil { - return fmt.Errorf("scheduler queries error: %w", err) - } - if err := q.doRegistryQueries(ctx, rng, height); err != nil { - return fmt.Errorf("registry queries error: %w", err) - } - if err := q.doStakingQueries(ctx, rng, height); err != nil { - return fmt.Errorf("staking queries error: %w", err) - } - if err := q.doRuntimeQueries(ctx, rng); err != nil { - return fmt.Errorf("runtime queries error: %w", err) + err := q.doQueries(loopCtx, rng) + cancel() + if err != nil { + return err } - q.logger.Debug("Queries done", - "height", height, - "height_latest", block.Height, - ) - select { case <-time.After(1 * time.Second): case <-gracefulExit.Done():