apollo_gateway: move get_state_reader_from_latest_block to be async

[itamar-starkware]

TL;DR

Made get_state_reader_from_latest_block an async function to improve performance and resource utilization.

What changed?

• Added the #[async_trait] attribute to the StateReaderFactory trait and its implementations
• Changed get_state_reader_from_latest_block from a synchronous to an async function
• Updated the SyncStateReaderFactory implementation to use await instead of block_on
• Modified StatefulTransactionValidatorFactory to handle the async state reader
• Updated tests to work with the new async interface

How to test?

• Run the existing test suite with cargo test
• Verify that test_instantiate_validator passes with the new async implementation
• Check that all implementations of StateReaderFactory correctly handle the async behavior

Why make this change?

Converting the state reader interface to async improves performance by avoiding blocking operations. This prevents thread pool exhaustion and allows for better resource utilization, especially during high-load scenarios. The change is part of a larger effort to make the codebase more asynchronous and efficient.

[itamar-starkware]

• apollo_gateway: move instantiate_validator to be async #10122 : 2 dependent PRs (#10123 , #10486 )
• apollo_gateway: move get_latest_block_info to be async #10121
• apollo_gateway: move get_block_info to be async #10120
• apollo_gateway: move get_state_reader_from_latest_block to be async #10119 👈 (View in Graphite)
• main-v0.14.1

This stack of pull requests is managed by Graphite. Learn more about stacking.

[reviewable-StarkWare]

This change is 

[TzahiTaub]

@TzahiTaub reviewed 4 of 6 files at r1, all commit messages.
Reviewable status: 4 of 6 files reviewed, all discussions resolved (waiting on @ArniStarkware)

[TzahiTaub]

@TzahiTaub reviewed 1 of 6 files at r1.
Reviewable status: 5 of 6 files reviewed, 1 unresolved discussion (waiting on @ArniStarkware)

---

crates/apollo_gateway/src/stateful_transaction_validator.rs line 61 at r1 (raw file):

        // TODO(yael 6/5/2024): consider storing the block_info as part of the
        // StatefulTransactionValidator and update it only once a new block is created.
        let state_reader = tokio::runtime::Handle::current()

I'm not sure what happens here, as we're not in an async context. See this. Does it work at all? If so, is because we assume the instantiate_validator function was called using a tokio::task::spawn_blocking somewhere in the stack?

Code quote:

tokio::runtime::Handle::current()

[itamar-starkware]

crates/apollo_gateway/src/stateful_transaction_validator.rs line 61 at r1 (raw file):

Previously, TzahiTaub (Tzahi) wrote…

I'm not sure what happens here, as we're not in an async context. See this. Does it work at all? If so, is because we assume the instantiate_validator function was called using a tokio::task::spawn_blocking somewhere in the stack?

I will insert runtime as arguemnt to the function

[itamar-starkware]

crates/apollo_gateway/src/stateful_transaction_validator_test.rs line 179 at r1 (raw file):

    let validator = tokio::task::block_in_place(|| {
        stateful_validator_factory.instantiate_validator(&mock_state_reader_factory)
    });

1. Why I don't run instantiate_validator without any tokio wrapper:

    let validator = 
        stateful_validator_factory
            .instantiate_validator(&mock_state_reader_factory, tokio::runtime::Handle::current());

Cause (runtime error):

Cannot start a runtime from within a runtime. This happens because a function (like `block_on`) attempted to block the current thread while the thread is being used to drive asynchronous tasks.

2. Why I don't use spawn_blocking:

    let validator = tokio::task::spawn_blocking(|| {
        stateful_validator_factory
            .instantiate_validator(&mock_state_reader_factory, tokio::runtime::Handle::current())
    })

Cause (compilation error):

`dyn stateful_transaction_validator::StatefulTransactionValidatorTrait` cannot be sent between threads safely
the trait `std::marker::Send` is not implemented for `dyn stateful_transaction_validator::StatefulTransactionValidatorTrait`
required for `std::ptr::Unique<dyn stateful_transaction_validator::StatefulTransactionValidatorTrait>` to implement `std::marker::Send`

3. Why I use multi-threaded tokio and not single threaded tokio:
   #[tokio::test]
   Cause (runtime error):
   can call blocking only when running on the multi-threaded runtime

This tokio wrapper will be removed in the next PR when instantiate_validator will become himself an async function.

[TzahiTaub]

@TzahiTaub reviewed 8 of 8 files at r2, all commit messages.
Reviewable status: all files reviewed, 1 unresolved discussion (waiting on @ArniStarkware and @itamar-starkware)

---

crates/apollo_gateway/src/stateful_transaction_validator_test.rs line 179 at r1 (raw file):

Previously, itamar-starkware wrote…

1. Why I don't run instantiate_validator without any tokio wrapper:

    let validator = 
        stateful_validator_factory
            .instantiate_validator(&mock_state_reader_factory, tokio::runtime::Handle::current());

Cause (runtime error):

Cannot start a runtime from within a runtime. This happens because a function (like `block_on`) attempted to block the current thread while the thread is being used to drive asynchronous tasks.

2. Why I don't use spawn_blocking:

    let validator = tokio::task::spawn_blocking(|| {
        stateful_validator_factory
            .instantiate_validator(&mock_state_reader_factory, tokio::runtime::Handle::current())
    })

Cause (compilation error):

`dyn stateful_transaction_validator::StatefulTransactionValidatorTrait` cannot be sent between threads safely
the trait `std::marker::Send` is not implemented for `dyn stateful_transaction_validator::StatefulTransactionValidatorTrait`
required for `std::ptr::Unique<dyn stateful_transaction_validator::StatefulTransactionValidatorTrait>` to implement `std::marker::Send`

3. Why I use multi-threaded tokio and not single threaded tokio:
   #[tokio::test]
   Cause (runtime error):
   can call blocking only when running on the multi-threaded runtime

This tokio wrapper will be removed in the next PR when instantiate_validator will become himself an async function.

2. is what I wanted as it is very similar to what we do in the code, but we don't return the validator as done here in the test, so it compiles. This is OK, but should be reverted in the following PRs when possible.

[ArniStarkware]

@ArniStarkware reviewed 1 of 2 files at r3.
Reviewable status: 7 of 8 files reviewed, 1 unresolved discussion (waiting on @itamar-starkware and @TzahiTaub)

[ArniStarkware]

@ArniStarkware reviewed 5 of 8 files at r2, 1 of 2 files at r3, all commit messages.
Reviewable status: all files reviewed, 1 unresolved discussion (waiting on @itamar-starkware and @TzahiTaub)

[ArniStarkware]

Reviewable status: all files reviewed, 1 unresolved discussion (waiting on @TzahiTaub)

---

crates/apollo_gateway/src/sync_state_reader.rs line 362 at r3 (raw file):

        let latest_block_number = self
            .runtime
            // TODO(guy.f): Do we want to count this as well?

I wonder what that means...

(This is a joke about how I don't understand this comment - I am in favor of deleting this line. This comment is non-blocking. Still - it would be nice to know.)

Code quote:

// TODO(guy.f): Do we want to count this as well?

[ArniStarkware]

Reviewable status: all files reviewed, 1 unresolved discussion (waiting on @itamar-starkware and @TzahiTaub)

---

crates/apollo_gateway/src/stateful_transaction_validator_test.rs line 179 at r1 (raw file):

Previously, TzahiTaub (Tzahi) wrote…

2. is what I wanted as it is very similar to what we do in the code, but we don't return the validator as done here in the test, so it compiles. This is OK, but should be reverted in the following PRs when possible.

Add a todo to delete this.
IIUC, it was deleted in #10122, but it would be nice to confirm.

[TzahiTaub]

Reviewable status: all files reviewed, 1 unresolved discussion (waiting on @ArniStarkware)

---

crates/apollo_gateway/src/sync_state_reader.rs line 362 at r3 (raw file):

Previously, ArniStarkware (Arnon Hod) wrote…

I wonder what that means...

(This is a joke about how I don't understand this comment - I am in favor of deleting this line. This comment is non-blocking. Still - it would be nice to know.)

Asked what it means. Please revive it with the following (Sahaks wants it to stay).

// TODO(guy.f): the call to `get_latest_block_number()` is not counted in the storage metrics as it is done prior to the creation of 
SharedStateSyncClientMetricWrapper, directly via the SharedStateSyncClient. Fix this and add this call to the metrics.

Also, please ask the TODO creator before removing one when possible.

[TzahiTaub]

@TzahiTaub reviewed 2 of 2 files at r3.
Reviewable status: all files reviewed, 1 unresolved discussion (waiting on @ArniStarkware)

[itamar-starkware]

crates/apollo_gateway/src/stateful_transaction_validator_test.rs line 179 at r1 (raw file):

Previously, ArniStarkware (Arnon Hod) wrote…

Add a todo to delete this.
IIUC, it was deleted in #10122, but it would be nice to confirm.

It is indeed deleted.
Anyway I added the Todo and remove it when instntiate_validator become async in the 4-th PR.

[itamar-starkware]

crates/apollo_gateway/src/sync_state_reader.rs line 362 at r3 (raw file):

Previously, TzahiTaub (Tzahi) wrote…

Asked what it means. Please revive it with the following (Sahaks wants it to stay).

// TODO(guy.f): the call to `get_latest_block_number()` is not counted in the storage metrics as it is done prior to the creation of 
SharedStateSyncClientMetricWrapper, directly via the SharedStateSyncClient. Fix this and add this call to the metrics.

Also, please ask the TODO creator before removing one when possible.

Done.

[TzahiTaub]

@TzahiTaub reviewed 2 of 2 files at r4, all commit messages.
Reviewable status: complete! all files reviewed, all discussions resolved (waiting on @itamar-starkware)