Skip to content

Commit 17de197

Browse files
committed
Drop VssStore runtime in block_in_place()
It fixes the problem when `VssStore` is dropped from within an asynchronous context: > Cannot drop a runtime in a context where blocking is not allowed. > This happens when a runtime is dropped from within an asynchronous context.
1 parent 5586b69 commit 17de197

File tree

1 file changed

+49
-31
lines changed

1 file changed

+49
-31
lines changed

src/io/vss_store.rs

Lines changed: 49 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,7 @@ type CustomRetryPolicy = FilteredRetryPolicy<
4141
pub struct VssStore {
4242
client: VssClient<CustomRetryPolicy>,
4343
store_id: String,
44-
runtime: Runtime,
44+
runtime: Option<Runtime>,
4545
storable_builder: StorableBuilder<RandEntropySource>,
4646
key_obfuscator: KeyObfuscator,
4747
}
@@ -51,7 +51,7 @@ impl VssStore {
5151
base_url: String, store_id: String, vss_seed: [u8; 32],
5252
header_provider: Arc<dyn VssHeaderProvider>,
5353
) -> io::Result<Self> {
54-
let runtime = tokio::runtime::Builder::new_multi_thread().enable_all().build()?;
54+
let runtime = Some(tokio::runtime::Builder::new_multi_thread().enable_all().build()?);
5555
let (data_encryption_key, obfuscation_master_key) =
5656
derive_data_encryption_and_obfuscation_keys(&vss_seed);
5757
let key_obfuscator = KeyObfuscator::new(obfuscation_master_key);
@@ -137,18 +137,19 @@ impl KVStore for VssStore {
137137
key: self.build_key(primary_namespace, secondary_namespace, key)?,
138138
};
139139

140-
let resp =
141-
tokio::task::block_in_place(|| self.runtime.block_on(self.client.get_object(&request)))
142-
.map_err(|e| {
143-
let msg = format!(
144-
"Failed to read from key {}/{}/{}: {}",
145-
primary_namespace, secondary_namespace, key, e
146-
);
147-
match e {
148-
VssError::NoSuchKeyError(..) => Error::new(ErrorKind::NotFound, msg),
149-
_ => Error::new(ErrorKind::Other, msg),
150-
}
151-
})?;
140+
let resp = tokio::task::block_in_place(|| {
141+
self.runtime.as_ref().unwrap().block_on(self.client.get_object(&request))
142+
})
143+
.map_err(|e| {
144+
let msg = format!(
145+
"Failed to read from key {}/{}/{}: {}",
146+
primary_namespace, secondary_namespace, key, e
147+
);
148+
match e {
149+
VssError::NoSuchKeyError(..) => Error::new(ErrorKind::NotFound, msg),
150+
_ => Error::new(ErrorKind::Other, msg),
151+
}
152+
})?;
152153
// unwrap safety: resp.value must be always present for a non-erroneous VSS response, otherwise
153154
// it is an API-violation which is converted to [`VssError::InternalServerError`] in [`VssClient`]
154155
let storable = Storable::decode(&resp.value.unwrap().value[..]).map_err(|e| {
@@ -179,14 +180,16 @@ impl KVStore for VssStore {
179180
delete_items: vec![],
180181
};
181182

182-
tokio::task::block_in_place(|| self.runtime.block_on(self.client.put_object(&request)))
183-
.map_err(|e| {
184-
let msg = format!(
185-
"Failed to write to key {}/{}/{}: {}",
186-
primary_namespace, secondary_namespace, key, e
187-
);
188-
Error::new(ErrorKind::Other, msg)
189-
})?;
183+
tokio::task::block_in_place(|| {
184+
self.runtime.as_ref().unwrap().block_on(self.client.put_object(&request))
185+
})
186+
.map_err(|e| {
187+
let msg = format!(
188+
"Failed to write to key {}/{}/{}: {}",
189+
primary_namespace, secondary_namespace, key, e
190+
);
191+
Error::new(ErrorKind::Other, msg)
192+
})?;
190193

191194
Ok(())
192195
}
@@ -204,22 +207,27 @@ impl KVStore for VssStore {
204207
}),
205208
};
206209

207-
tokio::task::block_in_place(|| self.runtime.block_on(self.client.delete_object(&request)))
208-
.map_err(|e| {
209-
let msg = format!(
210-
"Failed to delete key {}/{}/{}: {}",
211-
primary_namespace, secondary_namespace, key, e
212-
);
213-
Error::new(ErrorKind::Other, msg)
214-
})?;
210+
tokio::task::block_in_place(|| {
211+
self.runtime.as_ref().unwrap().block_on(self.client.delete_object(&request))
212+
})
213+
.map_err(|e| {
214+
let msg = format!(
215+
"Failed to delete key {}/{}/{}: {}",
216+
primary_namespace, secondary_namespace, key, e
217+
);
218+
Error::new(ErrorKind::Other, msg)
219+
})?;
215220
Ok(())
216221
}
217222

218223
fn list(&self, primary_namespace: &str, secondary_namespace: &str) -> io::Result<Vec<String>> {
219224
check_namespace_key_validity(primary_namespace, secondary_namespace, None, "list")?;
220225

221226
let keys = tokio::task::block_in_place(|| {
222-
self.runtime.block_on(self.list_all_keys(primary_namespace, secondary_namespace))
227+
self.runtime
228+
.as_ref()
229+
.unwrap()
230+
.block_on(self.list_all_keys(primary_namespace, secondary_namespace))
223231
})
224232
.map_err(|e| {
225233
let msg = format!(
@@ -233,6 +241,16 @@ impl KVStore for VssStore {
233241
}
234242
}
235243

244+
impl Drop for VssStore {
245+
fn drop(&mut self) {
246+
let runtime = self.runtime.take();
247+
// Fixes:
248+
// > Cannot drop a runtime in a context where blocking is not allowed.
249+
// > This happens when a runtime is dropped from within an asynchronous context.
250+
tokio::task::block_in_place(move || drop(runtime));
251+
}
252+
}
253+
236254
fn derive_data_encryption_and_obfuscation_keys(vss_seed: &[u8; 32]) -> ([u8; 32], [u8; 32]) {
237255
let hkdf = |initial_key_material: &[u8], salt: &[u8]| -> [u8; 32] {
238256
let mut engine = HmacEngine::<sha256::Hash>::new(salt);

0 commit comments

Comments
 (0)