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

Code before:
            pb_ctx: &ctx.context,
            start_ts: Some(ctx.txn_start_ts),
            allowed_in_flashback: ctx.allowed_in_flashback,
            ..Default::default()
        };
        // need to pass start_ts and ranges to check memory locks for replica read
        if need_check_locks_in_replica_read(&ctx.context) {
            for r in &ctx.ranges {
                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)
            }
        }

Review comment (bug):
No need feature gate again.

```suggestion
            // 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");
```

Code after:
            pb_ctx: &ctx.context,
            start_ts: Some(ctx.txn_start_ts),
            allowed_in_flashback: ctx.allowed_in_flashback,
            ..Default::default()
        };
        // need to pass start_ts and ranges to check memory locks for replica read
        if need_check_locks_in_replica_read(&ctx.context) {
            for r in &ctx.ranges {
                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.