Repository: tikv/tikv
File: src/coprocessor/endpoint.rs (Rust)

Code before:
                let start_key = txn_types::Key::from_raw(r.get_start());
                let end_key = txn_types::Key::from_raw(r.get_end());
                let mut key_range = kvrpcpb::KeyRange::default();
                key_range.set_start_key(start_key.into_encoded());
                key_range.set_end_key(end_key.into_encoded());
                snap_ctx.key_ranges.push(key_range);
            }
        }
        let snapshot_future = kv::in_memory_snapshot(engine, snap_ctx).map_err(Error::from);
        async move {
            #[cfg(feature = "failpoints")]
            {
                // Sleep here to simulate slow snapshot retrieval that causes timeout
                // The failpoint can be configured with a sleep duration, e.g.:
                // fail::cfg("coprocessor_async_in_memory_snapshot_timeout",
                // "sleep(500)").unwrap();
                fail_point!("coprocessor_async_in_memory_snapshot_timeout");
            }
            snapshot_future.await
        }
    }

    /// Gets a snapshot with timeout based on the request deadline.
    ///
    /// Safety: This function must be called within a context where a TLS engine
    /// exists.
    async unsafe fn get_snapshot_with_timeout(ctx: &ReqContext) -> Result<E::IMSnap> {
        let snapshot_future = with_tls_engine(|engine| Self::async_in_memory_snapshot(engine, ctx));
        let max_duration_to_get_snapshot = ctx.deadline.remaining_duration();
        if max_duration_to_get_snapshot.is_zero() {
            return Err(Error::DeadlineExceeded);
        }
        match async_timeout(snapshot_future, max_duration_to_get_snapshot).await {
            Ok(snapshot) => snapshot,
            Err(e) => {
                warn!("timeout when getting snapshot";
                    "region_id" => &ctx.context.get_region_id(),
                    "max_duration_to_get_snapshot" => ?max_duration_to_get_snapshot,
                    "err" => ?e,
                );
                Err(Error::DeadlineExceeded)
            }
        }
    }

    /// The real implementation of handling a unary request.
    ///
    /// It first retrieves a snapshot, then builds the `RequestHandler` over the
    /// snapshot and the given `handler_builder`. Finally, it calls the unary
    /// request interface of the `RequestHandler` to process the request and
    /// produce a result.

Review comment (refactor):
NIT: The default deadline (60s) is too long, maybe a shorter constant duration, e.g. 1s or 5s is better here.

Code after:
                let start_key = txn_types::Key::from_raw(r.get_start());
                let end_key = txn_types::Key::from_raw(r.get_end());
                let mut key_range = kvrpcpb::KeyRange::default();
                key_range.set_start_key(start_key.into_encoded());
                key_range.set_end_key(end_key.into_encoded());
                snap_ctx.key_ranges.push(key_range);
            }
        }
        let snapshot_future = kv::in_memory_snapshot(engine, snap_ctx).map_err(Error::from);
        async move {
            // Sleep here to simulate slow snapshot retrieval that causes timeout
            // The failpoint can be configured with a sleep duration, e.g.:
            // fail::cfg("coprocessor_async_in_memory_snapshot_timeout",
            // "sleep(500)").unwrap();
            fail_point!("coprocessor_async_in_memory_snapshot_timeout");
            snapshot_future.await
        }
    }

    /// Gets a snapshot with timeout based on the request deadline.
    ///
    /// Safety: This function must be called within a context where a TLS engine
    /// exists.
    async unsafe fn get_snapshot_with_timeout(ctx: &ReqContext) -> Result<E::IMSnap> {
        let snapshot_future = with_tls_engine(|engine| Self::async_in_memory_snapshot(engine, ctx));
        let max_duration_to_get_snapshot = ctx.deadline.remaining_duration();
        if max_duration_to_get_snapshot.is_zero() {
            return Err(Error::DeadlineExceeded);
        }
        match async_timeout(snapshot_future, max_duration_to_get_snapshot).await {
            Ok(snapshot) => snapshot,
            Err(e) => {
                warn!("timeout when getting snapshot";
                    "region_id" => &ctx.context.get_region_id(),
                    "max_duration_to_get_snapshot" => ?max_duration_to_get_snapshot,
                    "err" => ?e,
                );
                Err(Error::DeadlineExceeded)
            }
        }
    }

    /// The real implementation of handling a unary request.
    ///
    /// It first retrieves a snapshot, then builds the `RequestHandler` over the
    /// snapshot and the given `handler_builder`. Finally, it calls the unary
    /// request interface of the `RequestHandler` to process the request and
    /// produce a result.
    async fn handle_unary_request_impl(
        semaphore: Option<Arc<Semaphore>>,
        mut tracker: Box<Tracker<E>>,