Skip to content

Commit

Permalink
Refactor Store Api into client side and driver side
Browse files Browse the repository at this point in the history 
Refactors the Store api into the driver (backend) implementation
and a client Store/StoreLike api. Store & StoreLike have their
sizes known at compile-time. This enables us to add templates
to the client-side to make it easier to work with, for example,
we no longer need to Pin the store and we'll be able to add
things like `digest: impl Into<DigestInfo>` to make the callers
life much easier.

towards: TraceMachina#934
  • Loading branch information
@allada
allada committed May 28, 2024
1 parent 08e2f2e commit 26ed4fd
Show file tree
Hide file tree
Showing 48 changed files with 1,042 additions and 1,212 deletions.
4 changes: 2 additions & 2 deletions deployment-examples/docker-compose/Dockerfile
View file
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,7 @@ RUN apt-get update \
DEBIAN_FRONTEND=noninteractive \
apt-get install --no-install-recommends -y \
npm=8.5.1~ds-1 \
git=1:2.34.1-1ubuntu1.10 \
git=1:2.34.1-1ubuntu1.11 \
gcc=4:11.2.0-1ubuntu1 \
g++=4:11.2.0-1ubuntu1 \
python3=3.10.6-1~22.04 \
Expand All @@ -38,7 +38,7 @@ RUN apt-get update \
DEBIAN_FRONTEND=noninteractive \
apt-get install --no-install-recommends -y \
npm=6.14.4+ds-1ubuntu2 \
git=1:2.25.1-1ubuntu3.11 \
git=1:2.25.1-1ubuntu3.12 \
gcc=4:9.3.0-1ubuntu2 \
g++=4:9.3.0-1ubuntu2 \
python3=3.8.2-0ubuntu2 \
Expand Down
22 changes: 7 additions & 15 deletions nativelink-scheduler/src/cache_lookup_scheduler.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,6 @@
// limitations under the License.

use std::collections::HashMap;
use std::pin::Pin;
use std::sync::Arc;

use async_trait::async_trait;
Expand All @@ -30,7 +29,7 @@ use nativelink_util::action_messages::{
use nativelink_util::background_spawn;
use nativelink_util::common::DigestInfo;
use nativelink_util::digest_hasher::DigestHasherFunc;
use nativelink_util::store_trait::Store;
use nativelink_util::store_trait::{Store, StoreLike};
use parking_lot::{Mutex, MutexGuard};
use scopeguard::guard;
use tokio::select;
Expand All @@ -50,7 +49,7 @@ type CheckActions = HashMap<ActionInfoHashKey, Arc<watch::Sender<Arc<ActionState
pub struct CacheLookupScheduler {
/// A reference to the AC to find existing actions in.
/// To prevent unintended issues, this store should probably be a CompletenessCheckingStore.
ac_store: Arc<dyn Store>,
ac_store: Store,
/// The "real" scheduler to use to perform actions if they were not found
/// in the action cache.
action_scheduler: Arc<dyn ActionScheduler>,
Expand All @@ -59,14 +58,13 @@ pub struct CacheLookupScheduler {
}

async fn get_action_from_store(
ac_store: Pin<&dyn Store>,
ac_store: &Store,
action_digest: DigestInfo,
instance_name: String,
digest_function: DigestHasherFunc,
) -> Option<ProtoActionResult> {
// If we are a GrpcStore we shortcut here, as this is a special store.
let any_store = ac_store.inner_store(Some(action_digest)).as_any();
if let Some(grpc_store) = any_store.downcast_ref::<GrpcStore>() {
if let Some(grpc_store) = ac_store.downcast_ref::<GrpcStore>(Some(action_digest)) {
let action_result_request = GetActionResultRequest {
instance_name,
action_digest: Some(action_digest.into()),
Expand Down Expand Up @@ -103,10 +101,7 @@ fn subscribe_to_existing_action(
}

impl CacheLookupScheduler {
pub fn new(
ac_store: Arc<dyn Store>,
action_scheduler: Arc<dyn ActionScheduler>,
) -> Result<Self, Error> {
pub fn new(ac_store: Store, action_scheduler: Arc<dyn ActionScheduler>) -> Result<Self, Error> {
Ok(Self {
ac_store,
action_scheduler,
Expand Down Expand Up @@ -169,17 +164,14 @@ impl ActionScheduler for CacheLookupScheduler {
let action_digest = current_state.action_digest();
let instance_name = action_info.instance_name().clone();
if let Some(action_result) = get_action_from_store(
Pin::new(ac_store.as_ref()),
&ac_store,
*action_digest,
instance_name,
current_state.unique_qualifier.digest_function,
)
.await
{
match Pin::new(ac_store.clone().as_ref())
.has(*action_digest)
.await
{
match ac_store.has(*action_digest).await {
Ok(Some(_)) => {
Arc::make_mut(&mut current_state).stage =
ActionStage::CompletedFromCache(action_result);
Expand Down
13 changes: 6 additions & 7 deletions nativelink-scheduler/tests/cache_lookup_scheduler_test.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,6 @@
// limitations under the License.

use std::collections::HashMap;
use std::pin::Pin;
use std::sync::Arc;
use std::time::UNIX_EPOCH;

Expand All @@ -33,7 +32,7 @@ use nativelink_store::memory_store::MemoryStore;
use nativelink_util::action_messages::{ActionInfoHashKey, ActionResult, ActionStage, ActionState};
use nativelink_util::common::DigestInfo;
use nativelink_util::digest_hasher::DigestHasherFunc;
use nativelink_util::store_trait::Store;
use nativelink_util::store_trait::{Store, StoreLike};
use prost::Message;
use tokio::sync::watch;
use tokio::{self};
Expand All @@ -42,15 +41,15 @@ use utils::scheduler_utils::{make_base_action_info, INSTANCE_NAME};

struct TestContext {
mock_scheduler: Arc<MockActionScheduler>,
ac_store: Arc<dyn Store>,
ac_store: Store,
cache_scheduler: CacheLookupScheduler,
}

fn make_cache_scheduler() -> Result<TestContext, Error> {
let mock_scheduler = Arc::new(MockActionScheduler::new());
let ac_store = Arc::new(MemoryStore::new(
let ac_store = Store::new(Arc::new(MemoryStore::new(
&nativelink_config::stores::MemoryStore::default(),
));
)));
let cache_scheduler = CacheLookupScheduler::new(ac_store.clone(), mock_scheduler.clone())?;
Ok(TestContext {
mock_scheduler,
Expand Down Expand Up @@ -91,8 +90,8 @@ mod cache_lookup_scheduler_tests {
let context = make_cache_scheduler()?;
let action_info = make_base_action_info(UNIX_EPOCH);
let action_result = ProtoActionResult::from(ActionResult::default());
let store_pin = Pin::new(context.ac_store.as_ref());
store_pin
context
.ac_store
.update_oneshot(*action_info.digest(), action_result.encode_to_vec().into())
.await?;
let (_forward_watch_channel_tx, forward_watch_channel_rx) =
Expand Down
18 changes: 7 additions & 11 deletions nativelink-service/src/ac_server.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,8 +14,6 @@

use std::collections::HashMap;
use std::fmt::Debug;
use std::pin::Pin;
use std::sync::Arc;

use bytes::BytesMut;
use nativelink_config::cas_server::{AcStoreConfig, InstanceName};
Expand All @@ -31,14 +29,14 @@ use nativelink_store::grpc_store::GrpcStore;
use nativelink_store::store_manager::StoreManager;
use nativelink_util::common::DigestInfo;
use nativelink_util::digest_hasher::make_ctx_for_hash_func;
use nativelink_util::store_trait::Store;
use nativelink_util::store_trait::{Store, StoreLike};
use prost::Message;
use tonic::{Request, Response, Status};
use tracing::{error_span, event, instrument, Level};

#[derive(Clone)]
pub struct AcStoreInfo {
store: Arc<dyn Store>,
store: Store,
read_only: bool,
}

Expand Down Expand Up @@ -97,14 +95,12 @@ impl AcServer {
.try_into()?;

// If we are a GrpcStore we shortcut here, as this is a special store.
let any_store = store_info.store.inner_store(Some(digest)).as_any();
if let Some(grpc_store) = any_store.downcast_ref::<GrpcStore>() {
if let Some(grpc_store) = store_info.store.downcast_ref::<GrpcStore>(Some(digest)) {
return grpc_store.get_action_result(Request::new(request)).await;
}

Ok(Response::new(
get_and_decode_digest::<ActionResult>(Pin::new(store_info.store.as_ref()), &digest)
.await?,
get_and_decode_digest::<ActionResult>(&store_info.store, &digest).await?,
))
}

Expand Down Expand Up @@ -132,8 +128,7 @@ impl AcServer {
.try_into()?;

// If we are a GrpcStore we shortcut here, as this is a special store.
let any_store = store_info.store.inner_store(Some(digest)).as_any();
if let Some(grpc_store) = any_store.downcast_ref::<GrpcStore>() {
if let Some(grpc_store) = store_info.store.downcast_ref::<GrpcStore>(Some(digest)) {
return grpc_store.update_action_result(Request::new(request)).await;
}

Expand All @@ -146,7 +141,8 @@ impl AcServer {
.encode(&mut store_data)
.err_tip(|| "Provided ActionResult could not be serialized")?;

Pin::new(store_info.store.as_ref())
store_info
.store
.update_oneshot(digest, store_data.freeze())
.await
.err_tip(|| "Failed to update in action cache")?;
Expand Down
25 changes: 11 additions & 14 deletions nativelink-service/src/bytestream_server.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -44,7 +44,7 @@ use nativelink_util::digest_hasher::{
use nativelink_util::proto_stream_utils::WriteRequestStreamWrapper;
use nativelink_util::resource_info::ResourceInfo;
use nativelink_util::spawn;
use nativelink_util::store_trait::{Store, UploadSizeInfo};
use nativelink_util::store_trait::{Store, StoreLike, UploadSizeInfo};
use nativelink_util::task::JoinHandleDropGuard;
use parking_lot::Mutex;
use tokio::time::sleep;
Expand Down Expand Up @@ -153,7 +153,7 @@ type BytesWrittenAndIdleStream = (Arc<AtomicU64>, Option<IdleStream>);
type SleepFn = Arc<dyn Fn() -> BoxFuture<'static, ()> + Send + Sync>;

pub struct ByteStreamServer {
stores: HashMap<String, Arc<dyn Store>>,
stores: HashMap<String, Store>,
// Max number of bytes to send on each grpc stream chunk.
max_bytes_per_stream: usize,
active_uploads: Arc<Mutex<HashMap<String, BytesWrittenAndIdleStream>>>,
Expand Down Expand Up @@ -206,7 +206,7 @@ impl ByteStreamServer {
fn create_or_join_upload_stream(
&self,
uuid: String,
store: Arc<dyn Store>,
store: Store,
digest: DigestInfo,
) -> Result<ActiveStreamGuard<'_>, Error> {
let (uuid, bytes_received) = match self.active_uploads.lock().entry(uuid) {
Expand Down Expand Up @@ -236,7 +236,7 @@ impl ByteStreamServer {
let store_update_fut = Box::pin(async move {
// We need to wrap `Store::update()` in a another future because we need to capture
// `store` to ensure it's lifetime follows the future and not the caller.
Pin::new(store.as_ref())
store
// Bytestream always uses digest size as the actual byte size.
.update(
digest,
Expand All @@ -260,7 +260,7 @@ impl ByteStreamServer {

async fn inner_read(
&self,
store: Arc<dyn Store>,
store: Store,
digest: DigestInfo,
read_request: ReadRequest,
) -> Result<Response<ReadStream>, Error> {
Expand Down Expand Up @@ -289,7 +289,7 @@ impl ByteStreamServer {
maybe_get_part_result: None,
get_part_fut: Box::pin(async move {
store
.get_part_arc(digest, tx, read_request.read_offset as usize, read_limit)
.get_part(digest, tx, read_request.read_offset as usize, read_limit)
.await
}),
});
Expand Down Expand Up @@ -383,7 +383,7 @@ impl ByteStreamServer {
)]
async fn inner_write(
&self,
store: Arc<dyn Store>,
store: Store,
digest: DigestInfo,
stream: WriteRequestStreamWrapper<Streaming<WriteRequest>, Status>,
) -> Result<Response<WriteResponse>, Error> {
Expand Down Expand Up @@ -520,8 +520,7 @@ impl ByteStreamServer {
let digest = DigestInfo::try_new(resource_info.hash.as_ref(), resource_info.expected_size)?;

// If we are a GrpcStore we shortcut here, as this is a special store.
let any_store = store_clone.inner_store(Some(digest)).as_any();
if let Some(grpc_store) = any_store.downcast_ref::<GrpcStore>() {
if let Some(grpc_store) = store_clone.downcast_ref::<GrpcStore>(Some(digest)) {
return grpc_store
.query_write_status(Request::new(query_request.clone()))
.await;
Expand All @@ -544,7 +543,7 @@ impl ByteStreamServer {
}
}

let has_fut = Pin::new(store_clone.as_ref()).has(digest);
let has_fut = store_clone.has(digest);
let Some(item_size) = has_fut.await.err_tip(|| "Failed to call .has() on store")? else {
return Err(make_err!(Code::NotFound, "{}", "not found"));
};
Expand Down Expand Up @@ -583,8 +582,7 @@ impl ByteStream for ByteStreamServer {
let digest = DigestInfo::try_new(resource_info.hash.as_ref(), resource_info.expected_size)?;

// If we are a GrpcStore we shortcut here, as this is a special store.
let any_store = store.inner_store(Some(digest)).as_any();
if let Some(grpc_store) = any_store.downcast_ref::<GrpcStore>() {
if let Some(grpc_store) = store.downcast_ref::<GrpcStore>(Some(digest)) {
let stream = grpc_store.read(Request::new(read_request)).await?;
return Ok(Response::new(Box::pin(stream)));
}
Expand Down Expand Up @@ -640,8 +638,7 @@ impl ByteStream for ByteStreamServer {
.err_tip(|| "Invalid digest input in ByteStream::write")?;

// If we are a GrpcStore we shortcut here, as this is a special store.
let any_store = store.inner_store(Some(digest)).as_any();
if let Some(grpc_store) = any_store.downcast_ref::<GrpcStore>() {
if let Some(grpc_store) = store.downcast_ref::<GrpcStore>(Some(digest)) {
return grpc_store.write(stream).await.map_err(|e| e.into());
}

Expand Down
Loading

0 comments on commit 26ed4fd

Please sign in to comment.