From 1db3addf43d6006684ff378e21a16b4192defa2a Mon Sep 17 00:00:00 2001
From: Meyer Jacobs <meyer@fb.com>
Date: Fri, 30 Apr 2021 20:37:21 -0700
Subject: [PATCH] indexedlogdatastore: Add public API for for clients to batch
 queries without locking for each individually

Summary: Add a Read/Write Guard API to IndexedLogHgIdDataStore which allows client code outside the module to perform a series of reads and writes without locking for each individually.

Reviewed By: kulshrax

Differential Revision: D28075788

fbshipit-source-id: 2a65a426f443e1a421198ad8b4c610e4822574f7
---
 .../revisionstore/src/indexedlogdatastore.rs  | 63 ++++++++++++++++++-
 1 file changed, 62 insertions(+), 1 deletion(-)

diff --git a/eden/scm/lib/revisionstore/src/indexedlogdatastore.rs b/eden/scm/lib/revisionstore/src/indexedlogdatastore.rs
index 678734fa2d..98fc7330e9 100644
--- a/eden/scm/lib/revisionstore/src/indexedlogdatastore.rs
+++ b/eden/scm/lib/revisionstore/src/indexedlogdatastore.rs
@@ -15,7 +15,7 @@ use anyhow::{bail, ensure, Result};
 use byteorder::{BigEndian, ReadBytesExt, WriteBytesExt};
 use futures::{FutureExt, StreamExt};
 use minibytes::Bytes;
-use parking_lot::RwLock;
+use parking_lot::{RwLock, RwLockReadGuard, RwLockWriteGuard};
 use tokio::task::spawn_blocking;
 
 use configparser::{config::ConfigSet, convert::ByteCount};
@@ -52,6 +52,12 @@ pub struct IndexedLogHgIdDataStore {
     extstored_policy: ExtStoredPolicy,
 }
 
+pub struct IndexedLogHgIdDataStoreReadGuard<'a>(RwLockReadGuard<'a, IndexedLogHgIdDataStoreInner>);
+
+pub struct IndexedLogHgIdDataStoreWriteGuard<'a>(
+    RwLockWriteGuard<'a, IndexedLogHgIdDataStoreInner>,
+);
+
 #[derive(Clone, Debug)]
 pub struct Entry {
     key: Key,
@@ -286,6 +292,61 @@ impl IndexedLogHgIdDataStore {
         self.inner.write().log.flush()?;
         Ok(())
     }
+
+    pub fn read_lock<'a>(&'a self) -> IndexedLogHgIdDataStoreReadGuard<'a> {
+        IndexedLogHgIdDataStoreReadGuard(self.inner.read())
+    }
+
+    pub fn write_lock<'a>(&'a self) -> IndexedLogHgIdDataStoreWriteGuard<'a> {
+        IndexedLogHgIdDataStoreWriteGuard(self.inner.write())
+    }
+}
+
+impl<'a> IndexedLogHgIdDataStoreReadGuard<'a> {
+    /// Write an entry to the IndexedLog
+    ///
+    /// Like IndexedLogHgIdDataStore::get_entry, but uses the already-acquired read lock.
+    pub fn get_entry(&self, key: Key) -> Result<Option<Entry>> {
+        Ok(self.get_raw_entry(&key)?.map(|e| e.with_key(key)))
+    }
+
+    /// Attempt to read an Entry from IndexedLog, without overwriting the Key (return Key path may not match the request Key path)
+    ///
+    /// Like IndexedLogHgIdDataStore::get_raw_entry, but uses the already-acquired read lock.
+    fn get_raw_entry(&self, key: &Key) -> Result<Option<Entry>> {
+        Entry::from_log(key, &self.0.log)
+    }
+}
+
+impl<'a> IndexedLogHgIdDataStoreWriteGuard<'a> {
+    /// Write an entry to the IndexedLog
+    ///
+    /// Like IndexedLogHgIdDataStore::get_entry, but uses the already-acquired write lock.
+    pub fn get_entry(&self, key: Key) -> Result<Option<Entry>> {
+        Ok(self.get_raw_entry(&key)?.map(|e| e.with_key(key)))
+    }
+
+    /// Attempt to read an Entry from IndexedLog, without overwriting the Key (return Key path may not match the request Key path)
+    ///
+    /// Like IndexedLogHgIdDataStore::get_raw_entry, but uses the already-acquired write lock.
+    fn get_raw_entry(&self, key: &Key) -> Result<Option<Entry>> {
+        Entry::from_log(key, &self.0.log)
+    }
+
+    /// Write an entry to the IndexedLog
+    ///
+    /// Like IndexedLogHgIdDataStore::put_entry, but uses the already-acquired write lock.
+    pub fn put_entry(&mut self, entry: Entry) -> Result<()> {
+        entry.write_to_log(&mut self.0.log)
+    }
+
+    /// Flush the underlying IndexedLog
+    ///
+    /// Like IndexedLogHgIdDataStore::flush_log, but uses the already-acquired write lock.
+    pub fn flush_log(&mut self) -> Result<()> {
+        self.0.log.flush()?;
+        Ok(())
+    }
 }
 
 impl KeyedValue for Entry {