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

Code before:
                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.
    async fn handle_unary_request_impl(
        semaphore: Option<Arc<Semaphore>>,
        mut tracker: Box<Tracker<E>>,
        handler_builder: RequestHandlerBuilder<E::IMSnap>,

Review comment (suggestion):
I suggest use `tokio::time::timeout(max_duration_to_get_snapshot, snapshot_future).await` here, as its simpler and  from its implementation, in the fast path there is actually no extra cost(the timeout is lazily inited) since in most cases the snapshot_future should return a ready directly. But the impl in GLOBAL_TIMER_HANDLE while create a real timer on every call.

Code after:
                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>>,
        handler_builder: RequestHandlerBuilder<E::IMSnap>,
    ) -> Result<MemoryTraceGuard<coppb::Response>> {
        with_tls_tracker(|tracker1| {
            record_network_in_bytes(tracker1.metrics.grpc_req_size);