From ff56ee7413f973e4f70feaf71a027d926ee42341 Mon Sep 17 00:00:00 2001
From: Jason Volk <jason@zemos.net>
Date: Sun, 8 Dec 2024 05:45:15 +0000
Subject: [PATCH] add back the consume_budget on cache hits

Signed-off-by: Jason Volk <jason@zemos.net>
---
 src/database/map/get.rs | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/src/database/map/get.rs b/src/database/map/get.rs
index 04f5d0ae..79556656 100644
--- a/src/database/map/get.rs
+++ b/src/database/map/get.rs
@@ -2,8 +2,9 @@ use std::{convert::AsRef, fmt::Debug, io::Write, sync::Arc};
 
 use arrayvec::ArrayVec;
 use conduit::{err, implement, utils::result::MapExpect, Err, Result};
-use futures::{future, Future, FutureExt};
+use futures::{Future, FutureExt};
 use serde::Serialize;
+use tokio::task;
 
 use crate::{
 	keyval::KeyBuf,
@@ -63,7 +64,9 @@ where
 
 	let cached = self.get_cached(key);
 	if matches!(cached, Err(_) | Ok(Some(_))) {
-		return future::ready(cached.map_expect("data found in cache")).boxed();
+		return task::consume_budget()
+			.map(move |()| cached.map_expect("data found in cache"))
+			.boxed();
 	}
 
 	debug_assert!(matches!(cached, Ok(None)), "expected status Incomplete");