diff --git a/.gitignore b/.gitignore index c0f63c5..d972476 100644 --- a/.gitignore +++ b/.gitignore @@ -13,6 +13,9 @@ Cargo.lock # MSVC Windows builds of rustc generate these, which store debugging information *.pdb +# Coverage output +coverage/ + # Ignore IDE specific files .idea/* .vscode/* diff --git a/Cargo.toml b/Cargo.toml index 9c5d2b6..1d07805 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -7,6 +7,9 @@ members= [ "examples/simple", "examples/simple-multiple-instances" ] +exclude = [ + "maelstrom-tests" +] [workspace.dependencies.cargo-husky] diff --git a/Makefile b/Makefile new file mode 100644 index 0000000..55235b5 --- /dev/null +++ b/Makefile @@ -0,0 +1,118 @@ +# groupcache — Makefile +# Run `make help` to see all available targets. + +.PHONY: help build check test test-all clippy fmt lint bench doc clean \ + coverage coverage-html coverage-lcov coverage-summary \ + maelstrom-test maelstrom-baseline maelstrom-partitions + +# ── Build & Check ──────────────────────────────────────────────────── + +help: ## Show this help + @grep -E '^[a-zA-Z_-]+:.*?## .*$$' $(MAKEFILE_LIST) | sort | \ + awk 'BEGIN {FS = ":.*?## "}; {printf " \033[36m%-22s\033[0m %s\n", $$1, $$2}' + +build: ## Build all workspace crates + cargo build --workspace + +build-release: ## Build optimized release + cargo build --workspace --release + +check: ## Type-check without building + cargo check --workspace --all-targets + +# ── Testing ────────────────────────────────────────────────────────── + +test: ## Run all tests + cargo test --workspace + +test-lib: ## Run only library unit + integration tests + cargo test -p groupcache + +test-integration: ## Run only integration tests + cargo test -p groupcache --test integration_test + +test-metrics: ## Run only metrics tests + cargo test -p groupcache --test metrics_test + +test-bincode: ## Run tests with bincode feature enabled + cargo test -p groupcache --features bincode --lib --tests + +test-all: test test-bincode ## Run all tests including feature variants + +# ── Code Quality ───────────────────────────────────────────────────── + +clippy: ## Run clippy with warnings as errors + cargo clippy --workspace --all-targets -- -D warnings + +fmt: ## Format code + cargo fmt --all + +fmt-check: ## Check formatting without modifying + cargo fmt --all -- --check + +lint: clippy fmt-check ## Run all lints (clippy + format check) + +# ── Benchmarks ─────────────────────────────────────────────────────── + +bench: ## Run all benchmarks + cargo bench -p groupcache + +bench-codec: ## Run serialization benchmarks (msgpack vs bincode) + cargo bench -p groupcache --bench codec_bench + +bench-routing: ## Run routing contention benchmarks (RwLock vs ArcSwap) + cargo bench -p groupcache --bench routing_contention + +# ── Coverage ───────────────────────────────────────────────────────── + +coverage: coverage-summary ## Show coverage summary (alias) + +coverage-summary: ## Show per-file coverage summary (merged default + bincode runs) + cargo llvm-cov clean --workspace + cargo llvm-cov --no-report --workspace + cargo llvm-cov --no-report -p groupcache --features bincode --lib --tests + cargo llvm-cov report --summary-only + +coverage-text: ## Show line-by-line coverage in terminal + cargo llvm-cov clean --workspace + cargo llvm-cov --no-report --workspace + cargo llvm-cov --no-report -p groupcache --features bincode --lib --tests + cargo llvm-cov report --text + +coverage-html: ## Generate HTML coverage report and open in browser + cargo llvm-cov clean --workspace + cargo llvm-cov --no-report --workspace + cargo llvm-cov --no-report -p groupcache --features bincode --lib --tests + cargo llvm-cov report --html --open + +coverage-lcov: ## Generate LCOV report (for CI integration) + cargo llvm-cov clean --workspace + cargo llvm-cov --no-report --workspace + cargo llvm-cov --no-report -p groupcache --features bincode --lib --tests + cargo llvm-cov report --lcov --output-path lcov.info + @echo "Written to lcov.info" + +# ── Documentation ──────────────────────────────────────────────────── + +doc: ## Build and open documentation + cargo doc --workspace --no-deps --open + +# ── Maelstrom Convergence Tests ────────────────────────────────────── + +maelstrom-test: ## Run all Maelstrom convergence test scenarios + $(MAKE) -C maelstrom-tests test-maelstrom + +maelstrom-baseline: ## Run Maelstrom baseline (no faults) + $(MAKE) -C maelstrom-tests test-baseline + +maelstrom-partitions: ## Run Maelstrom partition test + $(MAKE) -C maelstrom-tests test-partitions + +maelstrom-results: ## View Maelstrom results in browser + $(MAKE) -C maelstrom-tests serve-results + +# ── Cleanup ────────────────────────────────────────────────────────── + +clean: ## Remove build artifacts + cargo clean + rm -f lcov.info diff --git a/examples/kubernetes-service-discovery/Cargo.toml b/examples/kubernetes-service-discovery/Cargo.toml index bee94dd..d80f153 100644 --- a/examples/kubernetes-service-discovery/Cargo.toml +++ b/examples/kubernetes-service-discovery/Cargo.toml @@ -7,7 +7,7 @@ publish = false [dependencies] # This pulls version from main branch so that docker build works (docker was confused by paths) #groupcache = { git = "https://github.com/Petroniuss/groupcache.git" } -groupcache = { path = "../../groupcache" } +groupcache = { path = "../../groupcache", package = "groupcache-ng", features = ["kubernetes"] } tonic = "0.14.2" axum = "0.8.0" @@ -25,5 +25,5 @@ axum-prometheus = "0.9.0" anyhow = "1" async-trait = "0.1" -kube = { version = "2.0.1", features = ["runtime", "derive"] } -k8s-openapi = { version = "0.26.0", features = ["latest"] } +# kube and k8s-openapi come through groupcache's "kubernetes" feature +kube = { version = "3", features = ["runtime"] } diff --git a/examples/kubernetes-service-discovery/src/k8s.rs b/examples/kubernetes-service-discovery/src/k8s.rs index a2334ef..da21b3b 100644 --- a/examples/kubernetes-service-discovery/src/k8s.rs +++ b/examples/kubernetes-service-discovery/src/k8s.rs @@ -1,61 +1,2 @@ -use async_trait::async_trait; -use groupcache::{GroupcachePeer, ServiceDiscovery}; -use k8s_openapi::api::core::v1::Pod; -use kube::api::ListParams; -use kube::{Api, Client}; -use std::collections::HashSet; -use std::error::Error; -use std::net::SocketAddr; - -pub struct Kubernetes { - api: Api, -} - -pub struct KubernetesBuilder { - client: Option, -} - -impl KubernetesBuilder { - pub fn build(self) -> Kubernetes { - Kubernetes { - api: Api::default_namespaced(self.client.unwrap()), - } - } - pub fn client(mut self, client: Client) -> Self { - self.client = Some(client); - self - } -} - -impl Kubernetes { - pub fn builder() -> KubernetesBuilder { - KubernetesBuilder { client: None } - } -} - -#[async_trait] -impl ServiceDiscovery for Kubernetes { - async fn pull_instances( - &self, - ) -> Result, Box> { - let pods_with_label_query = ListParams::default().labels("app=groupcache-powered-backend"); - Ok(self - .api - .list(&pods_with_label_query) - .await - .unwrap() - .into_iter() - .filter_map(|pod| { - let status = pod.status?; - let pod_ip = status.pod_ip?; - - let Ok(ip) = pod_ip.parse() else { - return None; - }; - - let addr = SocketAddr::new(ip, 3000); - Some(GroupcachePeer::from_socket(addr)) - }) - .collect()) - } -} +// Re-export from the library's built-in kubernetes discovery. +pub use groupcache::discovery::kubernetes::KubernetesDiscovery; diff --git a/examples/kubernetes-service-discovery/src/main.rs b/examples/kubernetes-service-discovery/src/main.rs index 062df5d..e930aad 100644 --- a/examples/kubernetes-service-discovery/src/main.rs +++ b/examples/kubernetes-service-discovery/src/main.rs @@ -3,7 +3,7 @@ mod cache; mod k8s; use crate::cache::CachedValue; -use crate::k8s::Kubernetes; +use crate::k8s::KubernetesDiscovery; use anyhow::Context; use anyhow::Result; use axum::extract::{Path, State}; @@ -59,8 +59,15 @@ async fn main() -> Result<()> { let client = Client::try_default().await?; // Configuring groupcache to use Kubernetes API server for peer auto-discovery. - let groupcache = Groupcache::builder(addr.into(), loader) - .service_discovery(Kubernetes::builder().client(client).build()) + let discovery = KubernetesDiscovery::builder() + .client(client) + .label_selector("app=groupcache-powered-backend") + .port(pod_port.parse()?) + .build() + .map_err(|e| anyhow::anyhow!("{}", e))?; + + let groupcache = Groupcache::builder(addr.into(), loader, groupcache::CancellationToken::new()) + .service_discovery(discovery) .build(); // Example axum app with endpoint to retrieve value from groupcache. diff --git a/examples/simple-multiple-instances/Cargo.toml b/examples/simple-multiple-instances/Cargo.toml index e3696a3..de01711 100644 --- a/examples/simple-multiple-instances/Cargo.toml +++ b/examples/simple-multiple-instances/Cargo.toml @@ -5,7 +5,7 @@ edition = "2021" publish = false [dependencies] -groupcache = { path = "../../groupcache" } +groupcache = { path = "../../groupcache", package = "groupcache-ng" } tonic = "0.14.2" tokio = { version = "1.34", features = ["full"] } diff --git a/examples/simple-multiple-instances/src/main.rs b/examples/simple-multiple-instances/src/main.rs index f8277cb..e37eacb 100644 --- a/examples/simple-multiple-instances/src/main.rs +++ b/examples/simple-multiple-instances/src/main.rs @@ -87,7 +87,7 @@ pub async fn spawn_groupcache_instance_on_addr(addr: SocketAddr) -> Result Result<()> { let loader = ComputeProtectedValue; // we crate groupcache with only a single peer - this process - let groupcache = Groupcache::builder(addr.into(), loader).build(); + let groupcache = Groupcache::builder(addr.into(), loader, groupcache::CancellationToken::new()).build(); // we make 3 concurrent requests for hot key let key = "some-hot-requested-key"; diff --git a/groupcache-pb/Cargo.toml b/groupcache-pb/Cargo.toml index e537127..fe378b1 100644 --- a/groupcache-pb/Cargo.toml +++ b/groupcache-pb/Cargo.toml @@ -1,21 +1,22 @@ [package] -name = "groupcache-pb" -version = "0.3.0" +name = "groupcache-ng-pb" +version = "0.4.0" edition = "2021" -authors = ["Patryk Wojtyczek"] +authors = ["Patryk Wojtyczek", "Henrik Johansson "] categories = ["caching", "web-programming", "concurrency", "asynchronous"] -description = "groupcache protocol buffers - internal crate" -homepage = "https://github.com/Petroniuss/groupcache" +description = "Protocol buffer definitions for groupcache-ng (internal crate)" +homepage = "https://github.com/dahankzter/groupcache" keywords = ["distributed", "cache", "shard", "memcached", "gRPC"] license = "MIT" readme = "../readme.md" -repository = "https://github.com/Petroniuss/groupcache" +repository = "https://github.com/dahankzter/groupcache" [dependencies] prost = "0.14.1" tonic = "0.14.2" tonic-prost = "0.14.2" tonic-prost-build = { version = "0.14.2" } +tokio-stream = "0.1" [build-dependencies] tonic-prost-build = { version = "0.14.2" } \ No newline at end of file diff --git a/groupcache-pb/build.rs b/groupcache-pb/build.rs index 145c5bb..b5a998b 100644 --- a/groupcache-pb/build.rs +++ b/groupcache-pb/build.rs @@ -1,20 +1,12 @@ fn main() -> Result<(), Box> { - // skip codegen if there hasn't been any updates. - if true { - return Ok(()); + // Generated code is checked in so that users don't need protoc installed. + // To regenerate after updating the .proto or bumping tonic/prost versions, + // change `false` to `true` below and run `cargo build -p groupcache-pb`. + if false { + tonic_prost_build::configure() + .out_dir("src/") + .compile_protos(&["protos/groupcache.proto"], &["protos/"])?; } - let current_dir = std::env::current_dir()?; - if !current_dir.ends_with("groupcache-pb") { - return Err(format!( - "must be run from the root of the crate, instead was {:#?}", - current_dir - ) - .into()); - } - - tonic_prost_build::configure() - .out_dir("src/") - .compile_protos(&["protos/groupcache.proto"], &["protos/"])?; Ok(()) } diff --git a/groupcache-pb/protos/groupcache.proto b/groupcache-pb/protos/groupcache.proto index cf7697c..4166f33 100644 --- a/groupcache-pb/protos/groupcache.proto +++ b/groupcache-pb/protos/groupcache.proto @@ -16,7 +16,14 @@ message RemoveRequest { message RemoveResponse {} +message WatchRequest {} + +message InvalidationEvent { + string key = 1; +} + service Groupcache { rpc Get(GetRequest) returns (GetResponse) {}; rpc Remove(RemoveRequest) returns (RemoveResponse) {}; + rpc WatchInvalidations(WatchRequest) returns (stream InvalidationEvent) {}; } diff --git a/groupcache-pb/src/groupcache_pb.rs b/groupcache-pb/src/groupcache_pb.rs index 321cd44..5d8e203 100644 --- a/groupcache-pb/src/groupcache_pb.rs +++ b/groupcache-pb/src/groupcache_pb.rs @@ -16,6 +16,13 @@ pub struct RemoveRequest { } #[derive(Clone, Copy, PartialEq, Eq, Hash, ::prost::Message)] pub struct RemoveResponse {} +#[derive(Clone, Copy, PartialEq, Eq, Hash, ::prost::Message)] +pub struct WatchRequest {} +#[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] +pub struct InvalidationEvent { + #[prost(string, tag = "1")] + pub key: ::prost::alloc::string::String, +} /// Generated client implementations. pub mod groupcache_client { #![allow( @@ -149,6 +156,30 @@ pub mod groupcache_client { .insert(GrpcMethod::new("groupcache_pb.Groupcache", "Remove")); self.inner.unary(req, path, codec).await } + pub async fn watch_invalidations( + &mut self, + request: impl tonic::IntoRequest, + ) -> std::result::Result< + tonic::Response>, + tonic::Status, + > { + self.inner + .ready() + .await + .map_err(|e| { + tonic::Status::unknown( + format!("Service was not ready: {}", e.into()), + ) + })?; + let codec = tonic_prost::ProstCodec::default(); + let path = http::uri::PathAndQuery::from_static( + "/groupcache_pb.Groupcache/WatchInvalidations", + ); + let mut req = request.into_request(); + req.extensions_mut() + .insert(GrpcMethod::new("groupcache_pb.Groupcache", "WatchInvalidations")); + self.inner.server_streaming(req, path, codec).await + } } } /// Generated server implementations. @@ -172,6 +203,19 @@ pub mod groupcache_server { &self, request: tonic::Request, ) -> std::result::Result, tonic::Status>; + /// Server streaming response type for the WatchInvalidations method. + type WatchInvalidationsStream: tonic::codegen::tokio_stream::Stream< + Item = std::result::Result, + > + + std::marker::Send + + 'static; + async fn watch_invalidations( + &self, + request: tonic::Request, + ) -> std::result::Result< + tonic::Response, + tonic::Status, + >; } #[derive(Debug)] pub struct GroupcacheServer { @@ -335,6 +379,51 @@ pub mod groupcache_server { }; Box::pin(fut) } + "/groupcache_pb.Groupcache/WatchInvalidations" => { + #[allow(non_camel_case_types)] + struct WatchInvalidationsSvc(pub Arc); + impl + tonic::server::ServerStreamingService + for WatchInvalidationsSvc { + type Response = super::InvalidationEvent; + type ResponseStream = T::WatchInvalidationsStream; + type Future = BoxFuture< + tonic::Response, + tonic::Status, + >; + fn call( + &mut self, + request: tonic::Request, + ) -> Self::Future { + let inner = Arc::clone(&self.0); + let fut = async move { + ::watch_invalidations(&inner, request).await + }; + Box::pin(fut) + } + } + let accept_compression_encodings = self.accept_compression_encodings; + let send_compression_encodings = self.send_compression_encodings; + let max_decoding_message_size = self.max_decoding_message_size; + let max_encoding_message_size = self.max_encoding_message_size; + let inner = self.inner.clone(); + let fut = async move { + let method = WatchInvalidationsSvc(inner); + let codec = tonic_prost::ProstCodec::default(); + let mut grpc = tonic::server::Grpc::new(codec) + .apply_compression_config( + accept_compression_encodings, + send_compression_encodings, + ) + .apply_max_message_size_config( + max_decoding_message_size, + max_encoding_message_size, + ); + let res = grpc.server_streaming(method, req).await; + Ok(res) + }; + Box::pin(fut) + } _ => { Box::pin(async move { let mut response = http::Response::new( diff --git a/groupcache-pb/src/lib.rs b/groupcache-pb/src/lib.rs index 9d62e3f..750bc24 100644 --- a/groupcache-pb/src/lib.rs +++ b/groupcache-pb/src/lib.rs @@ -7,7 +7,7 @@ mod groupcache_pb; pub use crate::groupcache_pb::groupcache_client::GroupcacheClient; pub use crate::groupcache_pb::groupcache_server::Groupcache; use crate::groupcache_pb::groupcache_server::GroupcacheServer as GroupcacheGRPCServer; -pub use crate::groupcache_pb::{GetRequest, GetResponse, RemoveRequest, RemoveResponse}; +pub use crate::groupcache_pb::{GetRequest, GetResponse, RemoveRequest, RemoveResponse, WatchRequest, InvalidationEvent}; /// gRPC server implementing groupcache GET to retrieve values. pub type GroupcacheServer = GroupcacheGRPCServer; diff --git a/groupcache/Cargo.toml b/groupcache/Cargo.toml index c68c725..13674e7 100644 --- a/groupcache/Cargo.toml +++ b/groupcache/Cargo.toml @@ -1,35 +1,49 @@ [package] -name = "groupcache" -version = "0.3.0" -authors = ["Patryk Wojtyczek"] +name = "groupcache-ng" +version = "0.4.0" +authors = ["Patryk Wojtyczek", "Henrik Johansson "] edition = "2021" categories = ["caching", "web-programming", "concurrency", "asynchronous"] description = """ -groupcache is a distributed caching and cache-filling library, \ -intended as a replacement for a pool of memcached nodes in many cases. \ -It shards by key to select which peer is responsible for that key.\ +A distributed caching and cache-filling library with near-realtime \ +streaming invalidation, lock-free routing, and built-in service discovery. \ +Fork of groupcache with added features — pending upstream merge.\ """ -homepage = "https://github.com/Petroniuss/groupcache" -keywords = ["distributed", "cache", "shard", "memcached", "gRPC"] +homepage = "https://github.com/dahankzter/groupcache" +keywords = ["distributed", "cache", "invalidation", "consistent-hashing", "gRPC"] license = "MIT" readme = "readme.md" -repository = "https://github.com/Petroniuss/groupcache" +repository = "https://github.com/dahankzter/groupcache" + +[lib] +name = "groupcache" + +[features] +bincode = ["dep:bincode"] +kubernetes = ["dep:kube", "dep:k8s-openapi"] +consul = ["dep:reqwest"] [dependencies] -groupcache-pb = { path = "../groupcache-pb", version = "0.3.0" } +groupcache-pb = { path = "../groupcache-pb", version = "0.4.0", package = "groupcache-ng-pb" } tonic = "0.14.2" hashring = "0.3.3" anyhow = "1.0" thiserror = "2.0" -axum = "0.8.0" -tokio = { version = "1.34" , features = ["rt"]} +arc-swap = "1" +tokio = { version = "1.34" , features = ["rt", "sync", "time"]} +tokio-stream = { version = "0.1", features = ["sync"] } +tokio-util = "0.7" serde = { version = "1.0", features = ["derive"] } rmp-serde = "1.1" +bincode = { version = "1", optional = true } async-trait = "0.1.74" singleflight-async = "0.2.0" moka = { version = "0.12.1", features = ["future"] } log = "0.4.20" metrics = "0.24.0" +kube = { version = "3", features = ["runtime", "derive"], optional = true } +k8s-openapi = { version = "0.27", features = ["latest"], optional = true } +reqwest = { version = "0.12", features = ["json"], default-features = false, optional = true } [dev-dependencies] cargo-husky = { workspace = true } @@ -37,3 +51,14 @@ pretty_assertions = "1.4.0" tokio-stream = { version = "0.1.14", features = ["net"] } tokio = { version = "1.34", features = ["time", "test-util", "rt", "macros"] } rstest = "0.26.1" +criterion = { version = "0.5", features = ["html_reports", "async_tokio"] } +parking_lot = "0.12" +bincode = "1" + +[[bench]] +name = "codec_bench" +harness = false + +[[bench]] +name = "routing_contention" +harness = false diff --git a/groupcache/benches/codec_bench.rs b/groupcache/benches/codec_bench.rs new file mode 100644 index 0000000..e265b6d --- /dev/null +++ b/groupcache/benches/codec_bench.rs @@ -0,0 +1,181 @@ +use criterion::{black_box, criterion_group, criterion_main, Criterion}; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Serialize, Deserialize, PartialEq)] +struct SmallValue { + id: String, + data: String, +} + +#[derive(Clone, Serialize, Deserialize, PartialEq)] +struct LargeValue { + id: String, + name: String, + description: String, + tags: Vec, + metadata: Vec<(String, String)>, +} + +fn small_value() -> SmallValue { + SmallValue { + id: "user-12345".into(), + data: "cached-payload".into(), + } +} + +fn large_value() -> LargeValue { + LargeValue { + id: "entity-99999".into(), + name: "Example Entity with a Reasonably Long Name".into(), + description: "A description that is representative of typical cached \ + values in a distributed system, containing enough data to show \ + serialization overhead differences between formats." + .into(), + tags: vec![ + "distributed".into(), + "cache".into(), + "performance".into(), + "benchmark".into(), + ], + metadata: vec![ + ("region".into(), "us-east-1".into()), + ("version".into(), "3.2.1".into()), + ("owner".into(), "team-platform".into()), + ], + } +} + +fn bench_serialize(c: &mut Criterion) { + let small = small_value(); + let large = large_value(); + + let mut group = c.benchmark_group("serialize"); + + group.bench_function("msgpack/small", |b| { + b.iter(|| rmp_serde::to_vec(black_box(&small)).unwrap()) + }); + + group.bench_function("bincode/small", |b| { + b.iter(|| bincode::serialize(black_box(&small)).unwrap()) + }); + + group.bench_function("msgpack/large", |b| { + b.iter(|| rmp_serde::to_vec(black_box(&large)).unwrap()) + }); + + group.bench_function("bincode/large", |b| { + b.iter(|| bincode::serialize(black_box(&large)).unwrap()) + }); + + group.finish(); +} + +fn bench_deserialize(c: &mut Criterion) { + let small = small_value(); + let large = large_value(); + + let small_mp = rmp_serde::to_vec(&small).unwrap(); + let small_bc = bincode::serialize(&small).unwrap(); + let large_mp = rmp_serde::to_vec(&large).unwrap(); + let large_bc = bincode::serialize(&large).unwrap(); + + let mut group = c.benchmark_group("deserialize"); + + group.bench_function("msgpack/small", |b| { + b.iter(|| rmp_serde::from_slice::(black_box(&small_mp)).unwrap()) + }); + + group.bench_function("bincode/small", |b| { + b.iter(|| bincode::deserialize::(black_box(&small_bc)).unwrap()) + }); + + group.bench_function("msgpack/large", |b| { + b.iter(|| rmp_serde::from_slice::(black_box(&large_mp)).unwrap()) + }); + + group.bench_function("bincode/large", |b| { + b.iter(|| bincode::deserialize::(black_box(&large_bc)).unwrap()) + }); + + group.finish(); +} + +fn bench_roundtrip(c: &mut Criterion) { + let small = small_value(); + let large = large_value(); + + let mut group = c.benchmark_group("roundtrip"); + + group.bench_function("msgpack/small", |b| { + b.iter(|| { + let bytes = rmp_serde::to_vec(black_box(&small)).unwrap(); + let _: SmallValue = rmp_serde::from_slice(&bytes).unwrap(); + }) + }); + + group.bench_function("bincode/small", |b| { + b.iter(|| { + let bytes = bincode::serialize(black_box(&small)).unwrap(); + let _: SmallValue = bincode::deserialize(&bytes).unwrap(); + }) + }); + + group.bench_function("msgpack/large", |b| { + b.iter(|| { + let bytes = rmp_serde::to_vec(black_box(&large)).unwrap(); + let _: LargeValue = rmp_serde::from_slice(&bytes).unwrap(); + }) + }); + + group.bench_function("bincode/large", |b| { + b.iter(|| { + let bytes = bincode::serialize(black_box(&large)).unwrap(); + let _: LargeValue = bincode::deserialize(&bytes).unwrap(); + }) + }); + + group.finish(); +} + +fn bench_wire_size(c: &mut Criterion) { + let small = small_value(); + let large = large_value(); + + let mut group = c.benchmark_group("wire_size"); + + // Not really a benchmark but a convenient way to report sizes + group.bench_function("msgpack/small", |b| { + let bytes = rmp_serde::to_vec(&small).unwrap(); + eprintln!(" msgpack small: {} bytes", bytes.len()); + b.iter(|| bytes.len()) + }); + + group.bench_function("bincode/small", |b| { + let bytes = bincode::serialize(&small).unwrap(); + eprintln!(" bincode small: {} bytes", bytes.len()); + b.iter(|| bytes.len()) + }); + + group.bench_function("msgpack/large", |b| { + let bytes = rmp_serde::to_vec(&large).unwrap(); + eprintln!(" msgpack large: {} bytes", bytes.len()); + b.iter(|| bytes.len()) + }); + + group.bench_function("bincode/large", |b| { + let bytes = bincode::serialize(&large).unwrap(); + eprintln!(" bincode large: {} bytes", bytes.len()); + b.iter(|| bytes.len()) + }); + + group.finish(); +} + +criterion_group!( + benches, + bench_serialize, + bench_deserialize, + bench_roundtrip, + bench_wire_size +); +criterion_main!(benches); diff --git a/groupcache/benches/routing_contention.rs b/groupcache/benches/routing_contention.rs new file mode 100644 index 0000000..74f7c03 --- /dev/null +++ b/groupcache/benches/routing_contention.rs @@ -0,0 +1,284 @@ +//! Benchmarks comparing routing state implementations under contention. +//! +//! Three approaches: +//! RwLock: readers acquire read lock, writers acquire write lock +//! ArcSwap+clone: readers do atomic load, writers clone-mutate-store +//! Double-buffer: readers do atomic load, writers mutate standby then swap +//! +//! All use the same Vec-based HashRing (best cache locality for reads). + +use arc_swap::ArcSwap; +use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion}; +use hashring::HashRing; +use parking_lot::{Mutex, RwLock}; +use std::sync::Arc; + +const VNODES_PER_PEER: usize = 40; + +// ── Shared ring type ──────────────────────────────────────────────── + +#[derive(Clone)] +struct Ring { + ring: HashRing, +} + +impl Ring { + fn new(num_peers: usize) -> Self { + let mut ring = HashRing::new(); + for peer_id in 0..num_peers { + for vnode_id in 0..VNODES_PER_PEER { + ring.add(format!("10.0.0.{}:8080_{}", peer_id, vnode_id)); + } + } + Self { ring } + } + + fn lookup(&self, key: &str) -> &str { + self.ring.get(&key).expect("ring not empty") + } + + fn add_vnode(&mut self, vnode: String) { + self.ring.add(vnode); + } + +} + +// ── Double-buffer wrapper (mirrors production DoubleBufferedRouting) ─ + +struct DoubleBuf { + active: ArcSwap, + standby: Mutex, +} + +impl DoubleBuf { + fn new(ring: Ring) -> Self { + Self { + active: ArcSwap::from_pointee(ring.clone()), + standby: Mutex::new(ring), + } + } + + fn load(&self) -> arc_swap::Guard> { + self.active.load() + } + + fn mutate(&self, vnode: &str) { + let mut standby = self.standby.lock(); + standby.add_vnode(vnode.to_string()); + + let old_active = self.active.swap(Arc::new(standby.clone())); + let mut caught_up = Arc::try_unwrap(old_active).unwrap_or_else(|arc| (*arc).clone()); + caught_up.add_vnode(vnode.to_string()); + *standby = caught_up; + } +} + +// ── Benchmarks ────────────────────────────────────────────────────── + +fn bench_read_scaling(c: &mut Criterion) { + let rt = tokio::runtime::Builder::new_multi_thread() + .worker_threads(16) + .build() + .unwrap(); + + let ring = Ring::new(10); + let keys: Vec = (0..1000).map(|i| format!("cache-key-{}", i)).collect(); + + let mut group = c.benchmark_group("routing_lookup"); + group.sample_size(50); + + for num_readers in [1, 2, 4, 8, 16] { + // RwLock + let state = Arc::new(RwLock::new(ring.clone())); + group.bench_with_input(BenchmarkId::new("rwlock", num_readers), &num_readers, |b, &n| { + b.to_async(&rt).iter(|| { + let state = state.clone(); + let keys = keys.clone(); + async move { + let mut handles = Vec::with_capacity(n); + for rid in 0..n { + let s = state.clone(); + let k = keys.clone(); + handles.push(tokio::spawn(async move { + for i in 0..1000 { + let key = &k[(rid * 100 + i) % k.len()]; + std::hint::black_box(s.read().lookup(key)); + } + })); + } + for h in handles { h.await.unwrap(); } + } + }); + }); + + // ArcSwap (read path same as double-buffer) + let state = Arc::new(ArcSwap::from_pointee(ring.clone())); + group.bench_with_input(BenchmarkId::new("arcswap", num_readers), &num_readers, |b, &n| { + b.to_async(&rt).iter(|| { + let state = state.clone(); + let keys = keys.clone(); + async move { + let mut handles = Vec::with_capacity(n); + for rid in 0..n { + let s = state.clone(); + let k = keys.clone(); + handles.push(tokio::spawn(async move { + for i in 0..1000 { + let key = &k[(rid * 100 + i) % k.len()]; + std::hint::black_box(s.load().lookup(key)); + } + })); + } + for h in handles { h.await.unwrap(); } + } + }); + }); + } + group.finish(); +} + +fn bench_mixed_matrix(c: &mut Criterion) { + let rt = tokio::runtime::Builder::new_multi_thread() + .worker_threads(20) + .build() + .unwrap(); + + let ring = Ring::new(10); + let keys: Vec = (0..1000).map(|i| format!("cache-key-{}", i)).collect(); + + let mut group = c.benchmark_group("routing_mixed"); + group.sample_size(30); + + for num_readers in [1, 4, 8, 16] { + for num_writers in [0, 1, 4] { + let label = format!("{}r_{}w", num_readers, num_writers); + + // ── RwLock ── + let state = Arc::new(RwLock::new(ring.clone())); + group.bench_function(format!("rwlock/{}", label), |b| { + b.to_async(&rt).iter(|| { + let state = state.clone(); + let keys = keys.clone(); + async move { + let mut handles = Vec::with_capacity(num_readers + num_writers); + for rid in 0..num_readers { + let s = state.clone(); + let k = keys.clone(); + handles.push(tokio::spawn(async move { + for i in 0..1000 { + let key = &k[(rid * 100 + i) % k.len()]; + std::hint::black_box(s.read().lookup(key)); + } + })); + } + for wid in 0..num_writers { + let s = state.clone(); + handles.push(tokio::spawn(async move { + for i in 0..10 { + { + let vnode = format!("10.0.0.{}:8080_{}", 50 + wid, i); + s.write().add_vnode(vnode); + } + tokio::task::yield_now().await; + } + })); + } + for h in handles { h.await.unwrap(); } + } + }); + }); + + // ── ArcSwap + clone ── + let state = Arc::new(ArcSwap::from_pointee(ring.clone())); + group.bench_function(format!("arcswap_clone/{}", label), |b| { + b.to_async(&rt).iter(|| { + let state = state.clone(); + let keys = keys.clone(); + async move { + let mut handles = Vec::with_capacity(num_readers + num_writers); + for rid in 0..num_readers { + let s = state.clone(); + let k = keys.clone(); + handles.push(tokio::spawn(async move { + for i in 0..1000 { + let key = &k[(rid * 100 + i) % k.len()]; + std::hint::black_box(s.load().lookup(key)); + } + })); + } + for wid in 0..num_writers { + let s = state.clone(); + handles.push(tokio::spawn(async move { + for i in 0..10 { + let mut new = (**s.load()).clone(); + new.add_vnode(format!("10.0.0.{}:8080_{}", 50 + wid, i)); + s.store(Arc::new(new)); + tokio::task::yield_now().await; + } + })); + } + for h in handles { h.await.unwrap(); } + } + }); + }); + + // ── Double-buffer ── + let state = Arc::new(DoubleBuf::new(ring.clone())); + group.bench_function(format!("double_buf/{}", label), |b| { + b.to_async(&rt).iter(|| { + let state = state.clone(); + let keys = keys.clone(); + async move { + let mut handles = Vec::with_capacity(num_readers + num_writers); + for rid in 0..num_readers { + let s = state.clone(); + let k = keys.clone(); + handles.push(tokio::spawn(async move { + for i in 0..1000 { + let key = &k[(rid * 100 + i) % k.len()]; + std::hint::black_box(s.load().lookup(key)); + } + })); + } + for wid in 0..num_writers { + let s = state.clone(); + handles.push(tokio::spawn(async move { + for i in 0..10 { + s.mutate(&format!("10.0.0.{}:8080_{}", 50 + wid, i)); + tokio::task::yield_now().await; + } + })); + } + for h in handles { h.await.unwrap(); } + } + }); + }); + } + } + group.finish(); +} + +fn bench_clone_and_lookup_by_cluster_size(c: &mut Criterion) { + let mut group = c.benchmark_group("ring_by_cluster_size"); + + for num_peers in [5, 10, 50, 128, 256] { + let ring = Ring::new(num_peers); + let keys: Vec = (0..1000).map(|i| format!("cache-key-{}", i)).collect(); + + group.bench_function(format!("clone/{}nodes", num_peers), |b| { + b.iter(|| std::hint::black_box(ring.clone())) + }); + + group.bench_function(format!("lookup_1000/{}nodes", num_peers), |b| { + b.iter(|| { + for key in &keys { + std::hint::black_box(ring.lookup(key)); + } + }) + }); + } + group.finish(); +} + +criterion_group!(benches, bench_read_scaling, bench_mixed_matrix, bench_clone_and_lookup_by_cluster_size); +criterion_main!(benches); diff --git a/groupcache/src/codec.rs b/groupcache/src/codec.rs new file mode 100644 index 0000000..03ec64a --- /dev/null +++ b/groupcache/src/codec.rs @@ -0,0 +1,31 @@ +//! Serialization codec for peer-to-peer value transfer. +//! +//! By default, values are serialized using MessagePack (rmp-serde). +//! Enable the `bincode` feature for faster serialization at the cost +//! of a less compact wire format. + +use serde::{Deserialize, Serialize}; + +/// Serialize a value to bytes for peer-to-peer transfer. +#[cfg(not(feature = "bincode"))] +pub(crate) fn serialize(value: &T) -> Result, Box> { + rmp_serde::to_vec(value).map_err(|e| Box::new(e) as _) +} + +/// Deserialize a value from bytes received from a peer. +#[cfg(not(feature = "bincode"))] +pub(crate) fn deserialize Deserialize<'a>>(bytes: &[u8]) -> Result> { + rmp_serde::from_slice(bytes).map_err(|e| Box::new(e) as _) +} + +/// Serialize a value to bytes for peer-to-peer transfer. +#[cfg(feature = "bincode")] +pub(crate) fn serialize(value: &T) -> Result, Box> { + bincode::serialize(value).map_err(|e| Box::new(e) as _) +} + +/// Deserialize a value from bytes received from a peer. +#[cfg(feature = "bincode")] +pub(crate) fn deserialize Deserialize<'a>>(bytes: &[u8]) -> Result> { + bincode::deserialize(bytes).map_err(|e| Box::new(e) as _) +} diff --git a/groupcache/src/discovery/consul.rs b/groupcache/src/discovery/consul.rs new file mode 100644 index 0000000..52b77fa --- /dev/null +++ b/groupcache/src/discovery/consul.rs @@ -0,0 +1,161 @@ +//! Consul-based service discovery for groupcache. +//! +//! Discovers healthy peers by querying Consul's +//! [`/v1/health/service`](https://developer.hashicorp.com/consul/api-docs/health#list-service-instances-for-a-service) +//! endpoint. Requires the `consul` feature. +//! +//! # Example +//! +//! ```no_run +//! use groupcache::discovery::consul::ConsulDiscovery; +//! +//! let discovery = ConsulDiscovery::builder() +//! .service_name("groupcache") +//! .port(8080) +//! .build(); +//! ``` + +use crate::groupcache::GroupcachePeer; +use crate::service_discovery::ServiceDiscovery; +use async_trait::async_trait; +use serde::Deserialize; +use std::collections::HashSet; +use std::error::Error; +use std::net::SocketAddr; +use std::time::Duration; + +/// Consul-based service discovery. +/// +/// Queries Consul's health API for healthy instances of a named service +/// and returns their addresses as groupcache peers. +pub struct ConsulDiscovery { + client: reqwest::Client, + url: String, + port: u16, + poll_interval: Duration, +} + +/// Builder for [`ConsulDiscovery`]. +pub struct ConsulDiscoveryBuilder { + consul_addr: String, + service_name: Option, + port: Option, + poll_interval: Option, + tags: Vec, +} + +impl ConsulDiscoveryBuilder { + /// Set the Consul service name to query. Required. + pub fn service_name(mut self, name: impl Into) -> Self { + self.service_name = Some(name.into()); + self + } + + /// Set the Consul HTTP address. Default: `http://localhost:8500`. + pub fn consul_addr(mut self, addr: impl Into) -> Self { + self.consul_addr = addr.into(); + self + } + + /// Set the groupcache port. If not set, uses the port from Consul's + /// service registration. + pub fn port(mut self, port: u16) -> Self { + self.port = Some(port); + self + } + + /// Filter by a Consul service tag. + pub fn tag(mut self, tag: impl Into) -> Self { + self.tags.push(tag.into()); + self + } + + /// Set the polling interval. Default: 10 seconds. + pub fn poll_interval(mut self, interval: Duration) -> Self { + self.poll_interval = Some(interval); + self + } + + /// Build the [`ConsulDiscovery`] instance. + pub fn build(self) -> ConsulDiscovery { + let service = self.service_name.expect("ConsulDiscovery requires service_name"); + let mut url = format!( + "{}/v1/health/service/{}?passing=true", + self.consul_addr.trim_end_matches('/'), + service, + ); + for tag in &self.tags { + url.push_str(&format!("&tag={}", tag)); + } + + ConsulDiscovery { + client: reqwest::Client::new(), + url, + port: self.port.unwrap_or(0), + poll_interval: self.poll_interval.unwrap_or(Duration::from_secs(10)), + } + } +} + +impl ConsulDiscovery { + /// Create a new builder. + pub fn builder() -> ConsulDiscoveryBuilder { + ConsulDiscoveryBuilder { + consul_addr: "http://localhost:8500".into(), + service_name: None, + port: None, + poll_interval: None, + tags: Vec::new(), + } + } +} + +// Consul health API response types (minimal) +#[derive(Deserialize)] +struct HealthServiceEntry { + #[serde(rename = "Service")] + service: ConsulService, +} + +#[derive(Deserialize)] +struct ConsulService { + #[serde(rename = "Address")] + address: String, + #[serde(rename = "Port")] + port: u16, +} + +#[async_trait] +impl ServiceDiscovery for ConsulDiscovery { + async fn pull_instances( + &self, + ) -> Result, Box> { + let entries: Vec = self + .client + .get(&self.url) + .send() + .await? + .json() + .await?; + + let peers = entries + .into_iter() + .filter_map(|entry| { + let ip = entry.service.address.parse().ok()?; + let port = if self.port > 0 { + self.port + } else { + entry.service.port + }; + let addr = SocketAddr::new(ip, port); + Some(GroupcachePeer::from_socket(addr)) + }) + .collect(); + + Ok(peers) + } + + fn interval(&self) -> Duration { + self.poll_interval + } +} diff --git a/groupcache/src/discovery/dns.rs b/groupcache/src/discovery/dns.rs new file mode 100644 index 0000000..b0cdb2b --- /dev/null +++ b/groupcache/src/discovery/dns.rs @@ -0,0 +1,106 @@ +//! DNS-based service discovery for groupcache. +//! +//! Resolves a DNS name to IP addresses and uses them as peers. +//! Works with headless Kubernetes services, SRV records, or any +//! DNS name that resolves to multiple A/AAAA records. +//! +//! No external dependencies required — uses `tokio::net::lookup_host`. +//! +//! # Example +//! +//! ```no_run +//! use groupcache::discovery::dns::DnsDiscovery; +//! +//! // Headless k8s service: resolves to all pod IPs +//! let discovery = DnsDiscovery::new("my-cache.default.svc.cluster.local", 8080); +//! ``` + +use crate::groupcache::GroupcachePeer; +use crate::service_discovery::ServiceDiscovery; +use async_trait::async_trait; +use std::collections::HashSet; +use std::error::Error; +use std::time::Duration; + +/// DNS-based service discovery. +/// +/// Resolves a hostname to IP addresses on each poll and returns +/// them as groupcache peers on the configured port. +pub struct DnsDiscovery { + host_port: String, + port: u16, + poll_interval: Duration, +} + +impl DnsDiscovery { + /// Create a new DNS discovery that resolves the given hostname. + /// + /// The `port` is the groupcache gRPC port that peers listen on. + pub fn new(hostname: impl Into, port: u16) -> Self { + let hostname = hostname.into(); + Self { + host_port: format!("{}:{}", hostname, port), + port, + poll_interval: Duration::from_secs(10), + } + } + + /// Set the polling interval. Default: 10 seconds. + pub fn with_interval(mut self, interval: Duration) -> Self { + self.poll_interval = interval; + self + } +} + +#[async_trait] +impl ServiceDiscovery for DnsDiscovery { + async fn pull_instances( + &self, + ) -> Result, Box> { + let addrs = tokio::net::lookup_host(&self.host_port).await?; + + let peers = addrs + .map(|mut addr| { + // Ensure we use our configured port (DNS might return + // a different port from SRV records) + addr.set_port(self.port); + GroupcachePeer::from_socket(addr) + }) + .collect(); + + Ok(peers) + } + + fn interval(&self) -> Duration { + self.poll_interval + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn dns_resolves_localhost() { + let discovery = DnsDiscovery::new("localhost", 8080); + let peers = discovery.pull_instances().await.unwrap(); + assert!(!peers.is_empty(), "localhost should resolve to at least one address"); + for peer in &peers { + assert_eq!(peer.socket.port(), 8080); + } + } + + #[tokio::test] + async fn dns_error_on_nonexistent_host() { + let discovery = DnsDiscovery::new("this-host-does-not-exist.invalid", 8080); + let result = discovery.pull_instances().await; + assert!(result.is_err()); + } + + #[test] + fn custom_interval() { + let discovery = DnsDiscovery::new("example.com", 8080) + .with_interval(Duration::from_secs(30)); + assert_eq!(discovery.interval(), Duration::from_secs(30)); + } +} diff --git a/groupcache/src/discovery/kubernetes.rs b/groupcache/src/discovery/kubernetes.rs new file mode 100644 index 0000000..61108fb --- /dev/null +++ b/groupcache/src/discovery/kubernetes.rs @@ -0,0 +1,153 @@ +//! Kubernetes service discovery for groupcache. +//! +//! Discovers peers by listing pods matching a label selector via the +//! Kubernetes API server. Requires the `kubernetes` feature. +//! +//! # Example +//! +//! ```no_run +//! use groupcache::discovery::kubernetes::KubernetesDiscovery; +//! +//! # async fn example() -> Result<(), Box> { +//! let client = kube::Client::try_default().await?; +//! let discovery = KubernetesDiscovery::builder() +//! .client(client) +//! .label_selector("app=my-cache") +//! .port(8080) +//! .build()?; +//! # Ok(()) +//! # } +//! ``` + +use crate::groupcache::GroupcachePeer; +use crate::service_discovery::ServiceDiscovery; +use async_trait::async_trait; +use k8s_openapi::api::core::v1::Pod; +use kube::api::ListParams; +use kube::{Api, Client}; +use std::collections::HashSet; +use std::error::Error; +use std::net::SocketAddr; +use std::time::Duration; + +/// Kubernetes-based service discovery for groupcache peers. +/// +/// Lists pods matching a label selector and extracts pod IPs to build +/// the peer set. Uses the Kubernetes API server (via the `kube` crate). +pub struct KubernetesDiscovery { + api: Api, + label_selector: String, + port: u16, + poll_interval: Duration, +} + +/// Builder for [`KubernetesDiscovery`]. +pub struct KubernetesDiscoveryBuilder { + client: Option, + namespace: Option, + label_selector: Option, + port: Option, + poll_interval: Option, +} + +impl KubernetesDiscoveryBuilder { + /// Set the Kubernetes client. Required. + pub fn client(mut self, client: Client) -> Self { + self.client = Some(client); + self + } + + /// Set a specific namespace. Defaults to the pod's own namespace. + pub fn namespace(mut self, namespace: impl Into) -> Self { + self.namespace = Some(namespace.into()); + self + } + + /// Set the label selector for discovering peer pods. Required. + /// + /// Example: `"app=my-groupcache-service"` + pub fn label_selector(mut self, selector: impl Into) -> Self { + self.label_selector = Some(selector.into()); + self + } + + /// Set the port that groupcache peers listen on. Required. + pub fn port(mut self, port: u16) -> Self { + self.port = Some(port); + self + } + + /// Set the polling interval for service discovery. Default: 10 seconds. + pub fn poll_interval(mut self, interval: Duration) -> Self { + self.poll_interval = Some(interval); + self + } + + /// Build the [`KubernetesDiscovery`] instance. + /// + /// # Errors + /// + /// Returns an error if `client`, `label_selector`, or `port` were not set. + pub fn build(self) -> Result> { + let client = self + .client + .ok_or("KubernetesDiscovery requires a kube::Client")?; + let label_selector = self + .label_selector + .ok_or("KubernetesDiscovery requires a label_selector")?; + let port = self + .port + .ok_or("KubernetesDiscovery requires a port")?; + + let api = match self.namespace { + Some(ns) => Api::namespaced(client, &ns), + None => Api::default_namespaced(client), + }; + + Ok(KubernetesDiscovery { + api, + label_selector, + port, + poll_interval: self.poll_interval.unwrap_or(Duration::from_secs(10)), + }) + } +} + +impl KubernetesDiscovery { + /// Create a new builder. + pub fn builder() -> KubernetesDiscoveryBuilder { + KubernetesDiscoveryBuilder { + client: None, + namespace: None, + label_selector: None, + port: None, + poll_interval: None, + } + } +} + +#[async_trait] +impl ServiceDiscovery for KubernetesDiscovery { + async fn pull_instances( + &self, + ) -> Result, Box> { + let params = ListParams::default().labels(&self.label_selector); + let pods = self.api.list(¶ms).await?; + + let peers = pods + .into_iter() + .filter_map(|pod| { + let pod_ip = pod.status?.pod_ip?; + let ip = pod_ip.parse().ok()?; + let addr = SocketAddr::new(ip, self.port); + Some(GroupcachePeer::from_socket(addr)) + }) + .collect(); + + Ok(peers) + } + + fn interval(&self) -> Duration { + self.poll_interval + } +} diff --git a/groupcache/src/discovery/mod.rs b/groupcache/src/discovery/mod.rs new file mode 100644 index 0000000..bb03cf3 --- /dev/null +++ b/groupcache/src/discovery/mod.rs @@ -0,0 +1,15 @@ +//! Built-in service discovery implementations. +//! +//! Enable via feature flags: +//! - `kubernetes` — discover peers via Kubernetes API server +//! - `consul` — discover peers via Consul health API +//! +//! DNS discovery is always available (no extra dependencies). + +pub mod dns; + +#[cfg(feature = "consul")] +pub mod consul; + +#[cfg(feature = "kubernetes")] +pub mod kubernetes; diff --git a/groupcache/src/errors.rs b/groupcache/src/errors.rs index 687d01e..435139c 100644 --- a/groupcache/src/errors.rs +++ b/groupcache/src/errors.rs @@ -1,12 +1,65 @@ -use rmp_serde::decode::Error; +use rmp_serde::decode::Error as RmpError; use std::sync::Arc; use tonic::Status; +/// Error type for groupcache operations. +/// +/// Variants are public so callers can match on error causes to distinguish +/// loader failures from network issues, deserialization problems, etc. #[derive(thiserror::Error, Debug)] -#[error(transparent)] -pub struct GroupcacheError { - #[from] - error: InternalGroupcacheError, +pub enum GroupcacheError { + /// The [`ValueLoader`](crate::ValueLoader) returned an error. + #[error("Loading error: '{0}'")] + Loading(Box), + + /// gRPC transport error when communicating with a peer. + #[error("Transport error: '{}'", .0.message())] + Transport(Status), + + /// Failed to deserialize a value received from a peer. + #[error("Deserialization error: '{0}'")] + Deserialization(RmpError), + + /// Failed to establish a connection to one or more peers. + #[error("Connection error: '{0}'")] + Connection(tonic::transport::Error), + + /// Multiple connection errors when updating the peer set. + #[error("Connection errors: {0:?}")] + ConnectionErrors(Vec), + + /// A peer returned an empty response. + #[error("Peer returned empty value for key '{0}'")] + EmptyResponse(String), + + /// Internal or unexpected error. + #[error(transparent)] + Internal(anyhow::Error), +} + +impl From for GroupcacheError { + fn from(err: InternalGroupcacheError) -> Self { + match err { + InternalGroupcacheError::LocalLoader(e) => GroupcacheError::Loading(e), + InternalGroupcacheError::Transport(e) => GroupcacheError::Transport(e), + InternalGroupcacheError::Rmp(e) => GroupcacheError::Deserialization(e), + InternalGroupcacheError::Connection(e) => GroupcacheError::Connection(e), + InternalGroupcacheError::EmptyResponse(key) => GroupcacheError::EmptyResponse(key), + InternalGroupcacheError::ConnectionErrors(errs) => { + GroupcacheError::ConnectionErrors(errs.into_iter().map(Into::into).collect()) + } + InternalGroupcacheError::Anyhow(e) => GroupcacheError::Internal(e), + InternalGroupcacheError::Deduped(DedupedGroupcacheError(arc)) => { + // Try to unwrap the Arc to preserve the original error variant. + // Falls back to string representation if other references exist + // (e.g. singleflight distributed the error to multiple callers). + match Arc::try_unwrap(arc) { + Ok(inner) => inner.into(), + Err(arc) => GroupcacheError::Internal(anyhow::anyhow!("{}", arc)), + } + } + } + } } #[derive(thiserror::Error, Debug, Clone)] @@ -22,7 +75,7 @@ pub(crate) enum InternalGroupcacheError { Transport(#[from] Status), #[error(transparent)] - Rmp(#[from] Error), + Rmp(#[from] RmpError), #[error(transparent)] Deduped(#[from] DedupedGroupcacheError), @@ -35,4 +88,51 @@ pub(crate) enum InternalGroupcacheError { #[error(transparent)] Anyhow(#[from] anyhow::Error), + + #[error("Peer returned empty value for key '{0}'")] + EmptyResponse(String), +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn transport_error_conversion() { + let status = Status::unavailable("peer down"); + let internal = InternalGroupcacheError::Transport(status); + let public: GroupcacheError = internal.into(); + assert!(matches!(public, GroupcacheError::Transport(_))); + } + + #[test] + fn rmp_deserialization_error_conversion() { + let rmp_err: RmpError = rmp_serde::from_slice::(&[0xFF]).unwrap_err(); + let internal = InternalGroupcacheError::Rmp(rmp_err); + let public: GroupcacheError = internal.into(); + assert!(matches!(public, GroupcacheError::Deserialization(_))); + } + + #[test] + fn empty_response_error_conversion() { + let internal = InternalGroupcacheError::EmptyResponse("some-key".to_string()); + let public: GroupcacheError = internal.into(); + assert!(matches!(public, GroupcacheError::EmptyResponse(_))); + } + + #[test] + fn anyhow_error_conversion() { + let internal = InternalGroupcacheError::Anyhow(anyhow::anyhow!("unexpected")); + let public: GroupcacheError = internal.into(); + assert!(matches!(public, GroupcacheError::Internal(_))); + } + + #[test] + fn deduped_error_fallback_when_arc_shared() { + let inner = Arc::new(InternalGroupcacheError::EmptyResponse("key".to_string())); + let _extra_ref = inner.clone(); // keep a second reference so try_unwrap fails + let internal = InternalGroupcacheError::Deduped(DedupedGroupcacheError(inner)); + let public: GroupcacheError = internal.into(); + assert!(matches!(public, GroupcacheError::Internal(_))); + } } diff --git a/groupcache/src/groupcache.rs b/groupcache/src/groupcache.rs index 838060d..adb0c06 100644 --- a/groupcache/src/groupcache.rs +++ b/groupcache/src/groupcache.rs @@ -9,6 +9,7 @@ use groupcache_pb::GroupcacheClient; use groupcache_pb::GroupcacheServer; use serde::{Deserialize, Serialize}; use std::collections::HashSet; +use std::fmt; use std::net::SocketAddr; use std::sync::Arc; use tonic::transport::Channel; @@ -40,11 +41,16 @@ impl Groupcache { /// In order to construct [`Groupcache`] application needs to provide: /// - [`GroupcachePeer`] - necessary for routing /// - [`ValueLoader`] implementation + /// - [`CancellationToken`](tokio_util::sync::CancellationToken) - for graceful shutdown + /// + /// When the token is cancelled, all background tasks (service discovery, + /// invalidation watchers) will stop gracefully. pub fn builder( me: GroupcachePeer, loader: impl ValueLoader + 'static, + cancel: tokio_util::sync::CancellationToken, ) -> GroupcacheBuilder { - GroupcacheBuilder::new(me, Box::new(loader)) + GroupcacheBuilder::new(me, Box::new(loader), cancel) } /// Provided a given `key` @@ -160,6 +166,14 @@ impl Groupcache { pub fn addr(&self) -> SocketAddr { self.0.addr() } + + /// Returns a snapshot of the instance's current health and cache state. + /// + /// The library reports state; the application decides policy. See + /// [`status::HealthState`] for the possible states. + pub fn status(&self) -> crate::status::Status { + self.0.status() + } } /// [ValueLoader] loads a value for a particular key - which can be potentially expensive. @@ -238,8 +252,26 @@ impl GroupcachePeer { } } +impl fmt::Display for GroupcachePeer { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{}", self.socket) + } +} + impl From for GroupcachePeer { fn from(value: SocketAddr) -> Self { Self { socket: value } } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn groupcache_peer_display() { + let addr: SocketAddr = "127.0.0.1:9090".parse().unwrap(); + let peer = GroupcachePeer::from_socket(addr); + assert_eq!(peer.to_string(), "127.0.0.1:9090"); + } +} diff --git a/groupcache/src/groupcache_builder.rs b/groupcache/src/groupcache_builder.rs index f36d86c..4484a0b 100644 --- a/groupcache/src/groupcache_builder.rs +++ b/groupcache/src/groupcache_builder.rs @@ -4,12 +4,20 @@ use crate::service_discovery::run_service_discovery; use crate::{Groupcache, GroupcacheInner, GroupcachePeer, ServiceDiscovery, ValueLoader}; use moka::future::Cache; use std::sync::Arc; +use tokio_util::sync::CancellationToken; use tonic::transport::Endpoint; +/// Type alias for the cache key used by groupcache. +/// +/// `Arc` is used instead of `String` to avoid heap allocations when +/// cloning keys across the main cache, hot cache, and singleflight group. +pub type CacheKey = Arc; + /// Allows to build groupcache instance with customized caches, timeouts etc pub struct GroupcacheBuilder { me: GroupcachePeer, loader: Box>, + cancel: CancellationToken, options: Options, } @@ -17,10 +25,12 @@ impl GroupcacheBuilder { pub(crate) fn new( me: GroupcachePeer, loader: Box + Sized + 'static>, + cancel: CancellationToken, ) -> Self { Self { me, loader, + cancel, options: Options::default(), } } @@ -31,9 +41,9 @@ impl GroupcacheBuilder { /// There is one owner of a given key in a set of peers /// forming hash ring and this peer stores values in main_cache. /// - /// By default, - /// [`moka`] cache is used and this setter allows to customize this cache (eviction policy, size etc.) - pub fn main_cache(mut self, main_cache: Cache) -> Self { + /// By default, main_cache stores up to 100k items with no TTL. + /// This setter allows to customize this cache (eviction policy, size etc.) + pub fn main_cache(mut self, main_cache: Cache) -> Self { self.options.main_cache = main_cache; self } @@ -46,7 +56,7 @@ impl GroupcacheBuilder { /// /// By default, hot_cache stores up to 10k items and expires after 30s. /// Depending on use_case you may either disable hot_cache or tweak time_to_live. - pub fn hot_cache(mut self, hot_cache: Cache) -> Self { + pub fn hot_cache(mut self, hot_cache: Cache) -> Self { self.options.hot_cache = hot_cache; self } @@ -62,9 +72,22 @@ impl GroupcacheBuilder { self } + /// Set the timeout for gRPC requests to peer nodes. + /// + /// This controls how long a `get()` or `remove()` call waits for a + /// response from the key's owner before failing. Default: 10 seconds. + /// + /// For datacenter deployments, a shorter timeout (e.g., 500ms) with + /// fast failover to local load is recommended. + pub fn grpc_timeout(mut self, timeout: std::time::Duration) -> Self { + self.options.grpc_timeout = timeout; + self + } + /// Allows to customize HTTP/2 channels for peer-to-peer connections. /// - /// By default, request timeout is set to 10 seconds. + /// For most use cases, [`grpc_timeout`](Self::grpc_timeout) is sufficient. + /// This method provides full control over the tonic `Endpoint`. pub fn grpc_endpoint_builder( mut self, builder: impl Fn(Endpoint) -> Endpoint + Send + Sync + 'static, @@ -73,6 +96,22 @@ impl GroupcacheBuilder { self } + /// Enable streaming invalidation. + /// + /// When enabled, key owners push invalidation events to all peers via + /// persistent gRPC streams. Peers clear their hot cache entries in + /// near-realtime instead of waiting for TTL expiration. + /// + /// This also enables: + /// - Hot cache flush on stream disconnect (prevents stale data after network issues) + /// - Selective cache eviction on membership changes (peers joining/leaving) + /// + /// By default, streaming invalidation is disabled for backwards compatibility. + pub fn enable_invalidation_streaming(mut self) -> Self { + self.options.invalidation.enabled = true; + self + } + /// Enable automatic pull-based service discovery via [`ServiceDiscovery`] /// /// By default, automatic service is disabled, useful when: @@ -88,14 +127,17 @@ impl GroupcacheBuilder { let cache = Groupcache(Arc::new(GroupcacheInner::new( self.me, self.loader, + self.cancel.clone(), self.options, ))); if let Some(service_discovery) = service_discovery { - tokio::spawn(run_service_discovery( + let handle = tokio::spawn(run_service_discovery( Arc::downgrade(&cache.0), service_discovery, + self.cancel.clone(), )); + cache.0.set_service_discovery_abort(handle.abort_handle()); } cache diff --git a/groupcache/src/groupcache_inner.rs b/groupcache/src/groupcache_inner.rs index e9ac1c9..9ca400c 100644 --- a/groupcache/src/groupcache_inner.rs +++ b/groupcache/src/groupcache_inner.rs @@ -3,9 +3,11 @@ use crate::errors::InternalGroupcacheError::Anyhow; use crate::errors::{DedupedGroupcacheError, GroupcacheError, InternalGroupcacheError}; use crate::groupcache::{GroupcachePeer, GroupcachePeerClient, ValueBounds, ValueLoader}; +use crate::invalidation::InvalidationManager; use crate::metrics::{ METRIC_GET_TOTAL, METRIC_LOCAL_CACHE_HIT_TOTAL, METRIC_LOCAL_LOAD_ERROR_TOTAL, METRIC_LOCAL_LOAD_TOTAL, METRIC_REMOTE_LOAD_ERROR, METRIC_REMOTE_LOAD_TOTAL, + METRIC_REMOVE_TOTAL, }; use crate::options::Options; use crate::routing::{GroupcachePeerWithClient, RoutingState}; @@ -17,24 +19,29 @@ use moka::future::Cache; use singleflight_async::SingleFlight; use std::collections::HashSet; use std::net::SocketAddr; -use std::sync::{Arc, RwLock}; -use tokio::task::JoinSet; +use arc_swap::ArcSwap; +use std::sync::{Arc, OnceLock}; +use tokio::task::{AbortHandle, JoinSet}; use tonic::transport::Endpoint; use tonic::IntoRequest; /// Core implementation of groupcache API. pub struct GroupcacheInner { - routing_state: Arc>, - single_flight_group: SingleFlight>, - main_cache: Cache, - hot_cache: Cache, + routing_state: Arc>, + single_flight_group: SingleFlight, Result>, + main_cache: Cache, Value>, + hot_cache: Cache, Value>, loader: Box>, config: Config, me: GroupcachePeer, + cancel: tokio_util::sync::CancellationToken, + service_discovery_abort: OnceLock, + pub(crate) invalidation: InvalidationManager, } struct Config { https: bool, + grpc_timeout: std::time::Duration, grpc_endpoint_builder: Arc Endpoint + Send + Sync + 'static>>, } @@ -42,9 +49,10 @@ impl GroupcacheInner { pub(crate) fn new( me: GroupcachePeer, loader: Box>, + cancel: tokio_util::sync::CancellationToken, options: Options, ) -> Self { - let routing_state = Arc::new(RwLock::new(RoutingState::with_local_peer(me))); + let routing_state = Arc::new(ArcSwap::from_pointee(RoutingState::with_local_peer(me))); let main_cache = options.main_cache; let hot_cache = options.hot_cache; @@ -53,9 +61,12 @@ impl GroupcacheInner { let config = Config { https: options.https, + grpc_timeout: options.grpc_timeout, grpc_endpoint_builder: Arc::new(options.grpc_endpoint_builder), }; + let invalidation = InvalidationManager::new(options.invalidation, cancel.clone()); + Self { routing_state, single_flight_group, @@ -64,6 +75,9 @@ impl GroupcacheInner { loader, me, config, + cancel, + service_discovery_abort: OnceLock::new(), + invalidation, } } @@ -88,7 +102,7 @@ impl GroupcacheInner { } let peer = { - let lock = self.routing_state.read().unwrap(); + let lock = self.routing_state.load(); lock.lookup_peer(key) }?; @@ -102,7 +116,7 @@ impl GroupcacheInner { peer: GroupcachePeerWithClient, ) -> Result { self.single_flight_group - .work(key.to_owned(), || async { + .work(Arc::from(key), || async { self.get_deduped(key, peer) .await .map_err(|e| DedupedGroupcacheError(Arc::new(e))) @@ -118,7 +132,7 @@ impl GroupcacheInner { ) -> Result { if peer.peer == self.me { let value = self.load_locally_instrumented(key).await?; - self.main_cache.insert(key.to_string(), value.clone()).await; + self.main_cache.insert(Arc::from(key), value.clone()).await; return Ok(value); } @@ -128,10 +142,14 @@ impl GroupcacheInner { let res = self.load_remotely_instrumented(key, &mut client).await; match res { Ok(value) => { - self.hot_cache.insert(key.to_string(), value.clone()).await; + self.hot_cache.insert(Arc::from(key), value.clone()).await; Ok(value) } - Err(_) => { + Err(err) => { + log::warn!( + "Remote load failed for key '{}' from peer {:?}, falling back to local load: {}", + key, peer.peer, err + ); let value = self.load_locally_instrumented(key).await?; Ok(value) } @@ -175,8 +193,11 @@ impl GroupcacheInner { .await?; let get_response = response.into_inner(); - let bytes = get_response.value.unwrap(); - let value = rmp_serde::from_read(bytes.as_slice())?; + let bytes = get_response + .value + .ok_or_else(|| InternalGroupcacheError::EmptyResponse(key.to_string()))?; + let value = crate::codec::deserialize(&bytes) + .map_err(|e| anyhow::anyhow!("failed to deserialize value from peer: {}", e))?; Ok(value) } @@ -185,15 +206,17 @@ impl GroupcacheInner { &self, key: &str, ) -> core::result::Result<(), InternalGroupcacheError> { + counter!(METRIC_REMOVE_TOTAL).increment(1); self.hot_cache.remove(key).await; let peer = { - let lock = self.routing_state.read().unwrap(); + let lock = self.routing_state.load(); lock.lookup_peer(key) }?; if peer.peer == self.me { self.main_cache.remove(key).await; + self.invalidation.broadcast(key); } else { let mut client = peer .client @@ -222,18 +245,33 @@ impl GroupcacheInner { } pub(crate) async fn add_peer(&self, peer: GroupcachePeer) -> Result<(), GroupcacheError> { - let contains_peer = { - let read_lock = self.routing_state.read().unwrap(); - read_lock.contains_peer(&peer) - }; - - if contains_peer { - return Ok(()); + // Quick check with read lock to avoid unnecessary connections. + { + let read_lock = self.routing_state.load(); + if read_lock.contains_peer(&peer) { + return Ok(()); + } } let (_, client) = self.connect(peer).await?; - let mut write_lock = self.routing_state.write().unwrap(); - write_lock.add_peer(peer, client); + + if self.invalidation.config().enabled { + self.invalidation + .spawn_watcher(peer, client.clone(), self.hot_cache.clone()); + } + + // Re-check to avoid adding a duplicate if another task connected concurrently. + { + let current = self.routing_state.load(); + if !current.contains_peer(&peer) { + let mut new_state = (**current).clone(); + new_state.add_peer(peer, client); + self.routing_state.store(Arc::new(new_state)); + } + } + + // After ring update: evict main cache keys we no longer own + self.evict_main_cache_keys_not_owned().await; Ok(()) } @@ -243,7 +281,7 @@ impl GroupcacheInner { updated_peers: HashSet, ) -> Result<(), GroupcacheError> { let current_peers: HashSet = { - let read_lock = self.routing_state.read().unwrap(); + let read_lock = self.routing_state.load(); read_lock.peers() }; @@ -257,12 +295,23 @@ impl GroupcacheInner { return Ok(()); } + // Before ring update: evict hot cache entries for keys owned by departing peers + for peer in &peers_to_remove { + self.evict_hot_cache_keys_owned_by(**peer).await; + } + let conn_errors = self.update_routing_table(new_connections_results, peers_to_remove); - conn_errors.is_empty().then(|| Ok(())).unwrap_or_else(|| { + + // After ring update: evict main cache keys we no longer own + self.evict_main_cache_keys_not_owned().await; + + if conn_errors.is_empty() { + Ok(()) + } else { Err(GroupcacheError::from( InternalGroupcacheError::ConnectionErrors(conn_errors), )) - }) + } } /// Updates routing table by adding new successful connections and removing old peers @@ -273,13 +322,17 @@ impl GroupcacheInner { >, peers_to_remove: Vec<&GroupcachePeer>, ) -> Vec { - let mut write_lock = self.routing_state.write().unwrap(); + let mut new_state = (**self.routing_state.load()).clone(); let mut connection_errors = Vec::new(); for result in connection_results { match result { Ok((peer, client)) => { - write_lock.add_peer(peer, client); + if self.invalidation.config().enabled { + self.invalidation + .spawn_watcher(peer, client.clone(), self.hot_cache.clone()); + } + new_state.add_peer(peer, client); } Err(e) => { connection_errors.push(e); @@ -288,9 +341,12 @@ impl GroupcacheInner { } for removed_peer in peers_to_remove { - write_lock.remove_peer(*removed_peer); + self.invalidation.cancel_watcher(removed_peer); + new_state.remove_peer(*removed_peer); } + self.routing_state.store(Arc::new(new_state)); + connection_errors } @@ -310,9 +366,10 @@ impl GroupcacheInner { for new_peer in peers_to_connect { let moved_peer = *new_peer; let https = self.config.https; + let timeout = self.config.grpc_timeout; let grpc_endpoint_builder = self.config.grpc_endpoint_builder.clone(); connection_task.spawn(async move { - GroupcacheInner::::connect_static(moved_peer, https, grpc_endpoint_builder) + GroupcacheInner::::connect_static(moved_peer, https, timeout, grpc_endpoint_builder) .await }); } @@ -336,6 +393,7 @@ impl GroupcacheInner { GroupcacheInner::::connect_static( peer, self.config.https, + self.config.grpc_timeout, self.config.grpc_endpoint_builder.clone(), ) .await @@ -344,6 +402,7 @@ impl GroupcacheInner { async fn connect_static( peer: GroupcachePeer, https: bool, + timeout: std::time::Duration, grpc_endpoint_builder: Arc Endpoint + Send + Sync + 'static>>, ) -> Result<(GroupcachePeer, GroupcachePeerClient), InternalGroupcacheError> { let socket = peer.socket; @@ -354,6 +413,7 @@ impl GroupcacheInner { }; let endpoint: Endpoint = peer_addr.try_into()?; + let endpoint = endpoint.timeout(timeout); let endpoint = grpc_endpoint_builder.as_ref()(endpoint); let client = GroupcacheClient::connect(endpoint).await?; Ok((peer, client)) @@ -361,7 +421,7 @@ impl GroupcacheInner { pub(crate) async fn remove_peer(&self, peer: GroupcachePeer) -> Result<(), GroupcacheError> { let contains_peer = { - let read_lock = self.routing_state.read().unwrap(); + let read_lock = self.routing_state.load(); read_lock.contains_peer(&peer) }; @@ -369,13 +429,105 @@ impl GroupcacheInner { return Ok(()); } - let mut write_lock = self.routing_state.write().unwrap(); - write_lock.remove_peer(peer); + self.invalidation.cancel_watcher(&peer); + + // Evict hot cache entries for keys owned by the departing peer + // (before ring update, while we can still identify them) + self.evict_hot_cache_keys_owned_by(peer).await; + + { + let mut new_state = (**self.routing_state.load()).clone(); + new_state.remove_peer(peer); + self.routing_state.store(Arc::new(new_state)); + } + + // After ring update: evict main cache keys we no longer own + // (peer removal can shift keys away from this node) + self.evict_main_cache_keys_not_owned().await; Ok(()) } + pub(crate) fn set_service_discovery_abort(&self, handle: AbortHandle) { + let _ = self.service_discovery_abort.set(handle); + } + pub(crate) fn addr(&self) -> SocketAddr { self.me.socket } + + pub(crate) fn status(&self) -> crate::status::Status { + use crate::status::{HealthState, Status}; + + let peer_count = self.routing_state.load().peers().len(); + + let health = if self.cancel.is_cancelled() { + HealthState::ShuttingDown + } else if peer_count == 0 { + HealthState::NoPeers + } else { + HealthState::Healthy + }; + + Status { + health, + peer_count, + main_cache_size: self.main_cache.entry_count(), + hot_cache_size: self.hot_cache.entry_count(), + } + } + + /// Evict hot cache entries for keys currently owned by the given peer. + /// Called before removing a peer from the ring, so ownership is still + /// queryable from the current ring state. + async fn evict_hot_cache_keys_owned_by(&self, peer: GroupcachePeer) { + let keys_to_evict: Vec> = { + let lock = self.routing_state.load(); + self.hot_cache + .iter() + .filter(|(key, _)| { + let key_str: &str = key.as_ref(); + lock.owner_of(key_str) == Some(peer) + }) + .map(|(key, _)| { + let inner: &Arc = &key; + Arc::clone(inner) + }) + .collect() + }; + for key in keys_to_evict { + self.hot_cache.remove(key.as_ref()).await; + } + } + + /// Evict main cache entries for keys this node no longer owns. + /// Called after a ring change (peer added) to clear keys that + /// moved to the new peer. + async fn evict_main_cache_keys_not_owned(&self) { + let keys_to_evict: Vec> = { + let lock = self.routing_state.load(); + self.main_cache + .iter() + .filter(|(key, _)| { + let key_str: &str = key.as_ref(); + lock.owner_of(key_str).is_some_and(|owner| owner != self.me) + }) + .map(|(key, _)| { + let inner: &Arc = &key; + Arc::clone(inner) + }) + .collect() + }; + for key in keys_to_evict { + self.main_cache.remove(key.as_ref()).await; + } + } +} + +impl Drop for GroupcacheInner { + fn drop(&mut self) { + if let Some(handle) = self.service_discovery_abort.get() { + handle.abort(); + } + } } diff --git a/groupcache/src/http.rs b/groupcache/src/http.rs index 4debcc1..6f344e7 100644 --- a/groupcache/src/http.rs +++ b/groupcache/src/http.rs @@ -4,8 +4,13 @@ use crate::groupcache::ValueBounds; use crate::metrics::METRIC_GET_SERVER_REQUESTS_TOTAL; use crate::GroupcacheInner; use async_trait::async_trait; -use groupcache_pb::{GetRequest, GetResponse, Groupcache, RemoveRequest, RemoveResponse}; +use groupcache_pb::{ + GetRequest, GetResponse, Groupcache, InvalidationEvent, RemoveRequest, RemoveResponse, + WatchRequest, +}; use metrics::counter; +use std::pin::Pin; +use tokio_stream::Stream; use tonic::{Request, Response, Status}; #[async_trait] @@ -16,8 +21,7 @@ impl Groupcache for GroupcacheInner { let payload = request.into_inner(); match self.get(&payload.key).await { Ok(value) => { - let result = rmp_serde::to_vec(&value); - match result { + match crate::codec::serialize(&value) { Ok(bytes) => Ok(Response::new(GetResponse { value: Some(bytes) })), Err(err) => Err(Status::internal(err.to_string())), } @@ -37,4 +41,15 @@ impl Groupcache for GroupcacheInner { Err(err) => Err(Status::internal(err.to_string())), } } + + type WatchInvalidationsStream = + Pin> + Send>>; + + async fn watch_invalidations( + &self, + _request: Request, + ) -> Result, Status> { + let stream = self.invalidation.subscribe_stream(); + Ok(Response::new(Box::pin(stream))) + } } diff --git a/groupcache/src/invalidation.rs b/groupcache/src/invalidation.rs new file mode 100644 index 0000000..e7b52b5 --- /dev/null +++ b/groupcache/src/invalidation.rs @@ -0,0 +1,186 @@ +use crate::groupcache::{GroupcachePeer, GroupcachePeerClient, ValueBounds}; +use crate::metrics::{ + METRIC_INVALIDATION_BROADCAST_TOTAL, METRIC_INVALIDATION_RECEIVED_TOTAL, + METRIC_INVALIDATION_STREAM_DISCONNECT_TOTAL, +}; +use groupcache_pb::{InvalidationEvent, WatchRequest}; +use log::{error, info, warn}; +use metrics::counter; +use moka::future::Cache; +use std::collections::HashMap; +use std::sync::{Arc, Mutex}; +use std::time::Duration; +use tokio::sync::broadcast; +use tokio::task::JoinHandle; +use tokio_stream::wrappers::BroadcastStream; +use tokio_stream::StreamExt; +use tonic::IntoRequest; + +/// Capacity of the broadcast channel. Events that arrive when all receivers are +/// this far behind are dropped (receivers will detect the lag and flush). +const BROADCAST_CHANNEL_CAPACITY: usize = 4096; + +/// Configuration for invalidation streaming. +#[derive(Clone)] +pub(crate) struct InvalidationConfig { + /// Whether streaming invalidation is enabled. + pub(crate) enabled: bool, + /// Base delay for exponential backoff on reconnection. + pub(crate) reconnect_base_delay: Duration, + /// Maximum delay for exponential backoff on reconnection. + pub(crate) reconnect_max_delay: Duration, +} + +impl Default for InvalidationConfig { + fn default() -> Self { + Self { + enabled: false, + reconnect_base_delay: Duration::from_millis(100), + reconnect_max_delay: Duration::from_secs(5), + } + } +} + +/// Manages invalidation event broadcasting and watcher streams. +pub(crate) struct InvalidationManager { + sender: broadcast::Sender, + config: InvalidationConfig, + cancel: tokio_util::sync::CancellationToken, + watchers: Mutex>>, +} + +impl InvalidationManager { + pub(crate) fn new(config: InvalidationConfig, cancel: tokio_util::sync::CancellationToken) -> Self { + let (sender, _) = broadcast::channel(BROADCAST_CHANNEL_CAPACITY); + Self { + cancel, + sender, + config, + watchers: Mutex::new(HashMap::new()), + } + } + + pub(crate) fn config(&self) -> &InvalidationConfig { + &self.config + } + + /// Broadcast an invalidation event for a key to all connected watchers. + /// Called by the key owner when a key is removed from main_cache. + pub(crate) fn broadcast(&self, key: &str) { + counter!(METRIC_INVALIDATION_BROADCAST_TOTAL).increment(1); + // send() returns Err only if there are no receivers -- that's fine, + // it means no peers are watching yet. + let _ = self.sender.send(key.to_string()); + } + + /// Create a new stream for a watcher (server-side of WatchInvalidations RPC). + /// Returns a stream of InvalidationEvent that tonic can serve. + pub(crate) fn subscribe_stream( + &self, + ) -> impl tokio_stream::Stream> { + let rx = self.sender.subscribe(); + BroadcastStream::new(rx).filter_map(|result| match result { + Ok(key) => Some(Ok(InvalidationEvent { key })), + Err(tokio_stream::wrappers::errors::BroadcastStreamRecvError::Lagged(n)) => { + warn!("Invalidation watcher lagged by {} events", n); + // Return empty-key event that the client interprets as "flush all" + Some(Ok(InvalidationEvent { + key: String::new(), + })) + } + }) + } + + /// Spawn a background task that watches a remote peer's invalidation stream + /// and removes keys from the local hot_cache. + pub(crate) fn spawn_watcher( + &self, + peer: GroupcachePeer, + mut client: GroupcachePeerClient, + hot_cache: Cache, V>, + ) { + let config = self.config.clone(); + let peer_socket = peer.socket; + let cancel = self.cancel.clone(); + + let handle = tokio::spawn(async move { + let mut backoff = config.reconnect_base_delay; + + loop { + if cancel.is_cancelled() { + break; + } + + match client + .watch_invalidations(WatchRequest {}.into_request()) + .await + { + Ok(response) => { + backoff = config.reconnect_base_delay; // reset on successful connect + let mut stream = response.into_inner(); + info!("Connected invalidation watcher to peer {:?}", peer_socket); + + while let Some(result) = stream.next().await { + match result { + Ok(event) => { + if event.key.is_empty() { + // Empty key = "flush all" signal (lagged or reconnect) + warn!( + "Received flush signal from peer {:?}, invalidating entire hot cache", + peer_socket + ); + hot_cache.invalidate_all(); + counter!(METRIC_INVALIDATION_STREAM_DISCONNECT_TOTAL) + .increment(1); + } else { + hot_cache.remove(event.key.as_str()).await; + counter!(METRIC_INVALIDATION_RECEIVED_TOTAL).increment(1); + } + } + Err(status) => { + warn!( + "Invalidation stream error from peer {:?}: {}", + peer_socket, status + ); + break; + } + } + } + } + Err(status) => { + error!( + "Failed to open invalidation stream to peer {:?}: {}", + peer_socket, status + ); + } + } + + // Stream ended or failed to connect -- flush hot cache and reconnect + warn!( + "Invalidation stream to peer {:?} disconnected, flushing hot cache", + peer_socket + ); + hot_cache.invalidate_all(); + counter!(METRIC_INVALIDATION_STREAM_DISCONNECT_TOTAL).increment(1); + + tokio::select! { + _ = tokio::time::sleep(backoff) => {} + _ = cancel.cancelled() => break, + } + backoff = (backoff * 2).min(config.reconnect_max_delay); + } + }); + + let mut watchers = self.watchers.lock().unwrap(); + if let Some(old_handle) = watchers.insert(peer, handle) { + old_handle.abort(); + } + } + + pub(crate) fn cancel_watcher(&self, peer: &GroupcachePeer) { + let mut watchers = self.watchers.lock().unwrap(); + if let Some(handle) = watchers.remove(peer) { + handle.abort(); + } + } +} diff --git a/groupcache/src/lib.rs b/groupcache/src/lib.rs index c11a948..a88c23c 100644 --- a/groupcache/src/lib.rs +++ b/groupcache/src/lib.rs @@ -1,20 +1,28 @@ #![doc = include_str!("../readme.md")] +mod codec; +pub mod discovery; mod errors; mod groupcache; mod groupcache_builder; mod groupcache_inner; mod http; +mod invalidation; pub mod metrics; mod options; mod routing; mod service_discovery; +pub mod status; pub use groupcache::{Groupcache, GroupcachePeer, ValueBounds, ValueLoader}; -pub use groupcache_builder::GroupcacheBuilder; +pub use groupcache_builder::{CacheKey, GroupcacheBuilder}; pub use groupcache_inner::GroupcacheInner; pub use groupcache_pb::GroupcacheServer; pub use service_discovery::ServiceDiscovery; /// we expose [`moka`](https://crates.io/crates/moka) since it's used in the public api of the library. pub use moka; + +/// Re-export [`CancellationToken`](tokio_util::sync::CancellationToken) since it's +/// required by [`Groupcache::builder`]. +pub use tokio_util::sync::CancellationToken; diff --git a/groupcache/src/metrics.rs b/groupcache/src/metrics.rs index 8d4b5a3..be7e565 100644 --- a/groupcache/src/metrics.rs +++ b/groupcache/src/metrics.rs @@ -1,31 +1,81 @@ -//! metrics module contains metric names along with description what each metric counts. +//! Metric names emitted by groupcache via the [`metrics`](https://docs.rs/metrics) facade. //! +//! Groupcache records counters using the `metrics` crate. To collect them, +//! install a metrics recorder in your application. Common options: //! -//! Metrics are exported via `metrics` create. -//! It's up to the application to ingest these metrics, some options: -//! - `metrics-exporter-prometheus`, -//! - `axum-prometheus`. +//! - [`metrics-exporter-prometheus`](https://docs.rs/metrics-exporter-prometheus) +//! - [`axum-prometheus`](https://docs.rs/axum-prometheus) //! -//! For meaning of each metric, see below. +//! All metric names are exported as public constants so applications can +//! reference them in dashboards, alerts, or tests. +//! +//! # Example +//! +//! ```no_run +//! // Install a Prometheus recorder at startup: +//! // let builder = metrics_exporter_prometheus::PrometheusBuilder::new(); +//! // builder.install().expect("failed to install recorder"); +//! +//! // Then groupcache metrics are automatically collected. +//! // Use the constants to build queries: +//! let hit_rate_query = format!( +//! "rate({}[5m]) / rate({}[5m])", +//! groupcache::metrics::LOCAL_CACHE_HIT_TOTAL, +//! groupcache::metrics::GET_TOTAL, +//! ); +//! ``` + +// ── Core operation counters ───────────────────────────────────────── + +/// Total number of get operations (cache hits + misses + remote fetches). +pub const GET_TOTAL: &str = "groupcache_get_total"; + +/// Total number of get requests received from remote peers via gRPC. +pub const GET_SERVER_REQUESTS_TOTAL: &str = "groupcache_get_server_requests_total"; + +/// Total number of local cache hits (main cache or hot cache). +/// No network call or loader invocation was needed. +pub const LOCAL_CACHE_HIT_TOTAL: &str = "groupcache_local_cache_hit_total"; + +/// Total number of times the [`ValueLoader`](crate::ValueLoader) was called. +pub const LOCAL_LOAD_TOTAL: &str = "groupcache_local_load_total"; + +/// Total number of [`ValueLoader`](crate::ValueLoader) failures. +pub const LOCAL_LOAD_ERROR_TOTAL: &str = "groupcache_local_load_errors"; + +/// Total number of remote fetches (gRPC calls to the key's owner). +pub const REMOTE_LOAD_TOTAL: &str = "groupcache_remote_load_total"; + +/// Total number of remote fetch failures. +pub const REMOTE_LOAD_ERROR_TOTAL: &str = "groupcache_remote_load_errors"; -/// Any GET request, including from peers. -pub(crate) const METRIC_GET_TOTAL: &str = "groupcache_get_total"; +/// Total number of remove (invalidation) operations initiated. +pub const REMOVE_TOTAL: &str = "groupcache_remove_total"; -/// GETs that came over the network from peers. -pub(crate) const METRIC_GET_SERVER_REQUESTS_TOTAL: &str = "groupcache_get_server_requests_total"; +// ── Streaming invalidation counters ───────────────────────────────── -/// Local cache hit (without going over the network or loading a value using [crate::groupcache::ValueLoader]). -pub(crate) const METRIC_LOCAL_CACHE_HIT_TOTAL: &str = "groupcache_local_cache_hit_total"; +/// Total number of invalidation events broadcast by this node (as key owner). +/// Only incremented when streaming invalidation is enabled. +pub const INVALIDATION_BROADCAST_TOTAL: &str = "groupcache_invalidation_broadcast_total"; -/// Total calls to [`crate::groupcache::ValueLoader::load`] -pub(crate) const METRIC_LOCAL_LOAD_TOTAL: &str = "groupcache_local_load_total"; +/// Total number of invalidation events received from remote peers via stream. +pub const INVALIDATION_RECEIVED_TOTAL: &str = "groupcache_invalidation_received_total"; -/// Total number of failures of [`crate::groupcache::ValueLoader::load`] -pub(crate) const METRIC_LOCAL_LOAD_ERROR_TOTAL: &str = "groupcache_local_load_errors"; +/// Total number of invalidation stream disconnects. +/// Each disconnect triggers a full hot cache flush as a safety measure. +pub const INVALIDATION_STREAM_DISCONNECT_TOTAL: &str = + "groupcache_invalidation_stream_disconnect_total"; -/// Total number of remote GETs: -/// - peer is not the owner and needs to make HTTP request to the owner for a given key. -pub(crate) const METRIC_REMOTE_LOAD_TOTAL: &str = "groupcache_remote_load_total"; +// ── Internal aliases (preserve old names for use inside the crate) ── -/// Total number of remote GET failures. -pub(crate) const METRIC_REMOTE_LOAD_ERROR: &str = "groupcache_remote_load_errors"; +pub(crate) use GET_TOTAL as METRIC_GET_TOTAL; +pub(crate) use GET_SERVER_REQUESTS_TOTAL as METRIC_GET_SERVER_REQUESTS_TOTAL; +pub(crate) use LOCAL_CACHE_HIT_TOTAL as METRIC_LOCAL_CACHE_HIT_TOTAL; +pub(crate) use LOCAL_LOAD_TOTAL as METRIC_LOCAL_LOAD_TOTAL; +pub(crate) use LOCAL_LOAD_ERROR_TOTAL as METRIC_LOCAL_LOAD_ERROR_TOTAL; +pub(crate) use REMOTE_LOAD_TOTAL as METRIC_REMOTE_LOAD_TOTAL; +pub(crate) use REMOTE_LOAD_ERROR_TOTAL as METRIC_REMOTE_LOAD_ERROR; +pub(crate) use REMOVE_TOTAL as METRIC_REMOVE_TOTAL; +pub(crate) use INVALIDATION_BROADCAST_TOTAL as METRIC_INVALIDATION_BROADCAST_TOTAL; +pub(crate) use INVALIDATION_RECEIVED_TOTAL as METRIC_INVALIDATION_RECEIVED_TOTAL; +pub(crate) use INVALIDATION_STREAM_DISCONNECT_TOTAL as METRIC_INVALIDATION_STREAM_DISCONNECT_TOTAL; diff --git a/groupcache/src/options.rs b/groupcache/src/options.rs index b977e7a..bf03208 100644 --- a/groupcache/src/options.rs +++ b/groupcache/src/options.rs @@ -1,39 +1,47 @@ +use crate::invalidation::InvalidationConfig; use crate::service_discovery::ServiceDiscovery; use crate::ValueBounds; use moka::future::Cache; +use std::sync::Arc; use std::time::Duration; use tonic::transport::Endpoint; +static DEFAULT_MAIN_CACHE_MAX_CAPACITY: u64 = 100_000; static DEFAULT_HOT_CACHE_MAX_CAPACITY: u64 = 10_000; static DEFAULT_HOT_CACHE_TIME_TO_LIVE: Duration = Duration::from_secs(30); static DEFAULT_GRPC_CLIENT_REQUEST_TIMEOUT: Duration = Duration::from_secs(10); pub(crate) struct Options { - pub(crate) main_cache: Cache, - pub(crate) hot_cache: Cache, + pub(crate) main_cache: Cache, Value>, + pub(crate) hot_cache: Cache, Value>, + pub(crate) grpc_timeout: Duration, pub(crate) grpc_endpoint_builder: Box Endpoint + Send + Sync + 'static>, pub(crate) https: bool, pub(crate) service_discovery: Option>, + pub(crate) invalidation: InvalidationConfig, } impl Default for Options { fn default() -> Self { - let main_cache = Cache::::builder().build(); + let main_cache = Cache::, Value>::builder() + .max_capacity(DEFAULT_MAIN_CACHE_MAX_CAPACITY) + .build(); - let hot_cache = Cache::::builder() + let hot_cache = Cache::, Value>::builder() .max_capacity(DEFAULT_HOT_CACHE_MAX_CAPACITY) .time_to_live(DEFAULT_HOT_CACHE_TIME_TO_LIVE) .build(); - let grpc_endpoint_builder = - Box::new(|e: Endpoint| e.timeout(DEFAULT_GRPC_CLIENT_REQUEST_TIMEOUT)); + let grpc_endpoint_builder = Box::new(|e: Endpoint| e); Self { main_cache, hot_cache, + grpc_timeout: DEFAULT_GRPC_CLIENT_REQUEST_TIMEOUT, grpc_endpoint_builder, https: false, service_discovery: None, + invalidation: InvalidationConfig::default(), } } } diff --git a/groupcache/src/routing.rs b/groupcache/src/routing.rs index d856b06..30015bb 100644 --- a/groupcache/src/routing.rs +++ b/groupcache/src/routing.rs @@ -2,10 +2,12 @@ use crate::groupcache::{GroupcachePeer, GroupcachePeerClient}; use anyhow::{Context, Result}; use hashring::HashRing; use std::collections::{HashMap, HashSet}; +use std::hash::{Hash, Hasher}; use std::net::SocketAddr; -static VNODES_PER_PEER: i32 = 40; +static VNODES_PER_PEER: u32 = 40; +#[derive(Clone)] pub(crate) struct RoutingState { peers: HashMap, ring: HashRing, @@ -45,6 +47,11 @@ impl RoutingState { Ok(GroupcachePeerWithClient { peer, client }) } + /// Returns the peer that owns the given key, or None if the ring is empty. + pub(crate) fn owner_of(&self, key: &str) -> Option { + self.ring.get(&key).map(|vnode| vnode.as_peer()) + } + fn peer_for_key(&self, key: &str) -> Result { let vnode = self .ring @@ -78,30 +85,30 @@ impl RoutingState { } } -#[derive(Debug, Clone, Hash, PartialEq, Eq)] +#[derive(Debug, Clone, PartialEq, Eq)] struct VNode { - addr_id: String, + addr: SocketAddr, + id: u32, } -impl VNode { - fn new(addr: SocketAddr, id: usize) -> Self { - Self { - addr_id: format!("{}_{}", addr, id), - } +impl Hash for VNode { + fn hash(&self, state: &mut H) { + // Preserve the same hash distribution as the original string-based format + // so that upgrading doesn't reshuffle the hash ring. + format!("{}_{}", self.addr, self.id).hash(state); } +} - fn vnodes_for_peer(peer: GroupcachePeer, num: i32) -> Vec { - let mut vnodes = Vec::new(); - for i in 0..num { - let vnode = VNode::new(peer.socket, i as usize); +impl VNode { + fn new(addr: SocketAddr, id: u32) -> Self { + Self { addr, id } + } - vnodes.push(vnode); - } - vnodes + fn vnodes_for_peer(peer: GroupcachePeer, num: u32) -> Vec { + (0..num).map(|i| VNode::new(peer.socket, i)).collect() } fn as_peer(&self) -> GroupcachePeer { - let addr = self.addr_id.split('_').next().unwrap().parse().unwrap(); - GroupcachePeer { socket: addr } + GroupcachePeer { socket: self.addr } } } diff --git a/groupcache/src/service_discovery.rs b/groupcache/src/service_discovery.rs index e8ff10d..47ef642 100644 --- a/groupcache/src/service_discovery.rs +++ b/groupcache/src/service_discovery.rs @@ -35,9 +35,14 @@ pub trait ServiceDiscovery: Send { pub(crate) async fn run_service_discovery( cache: Weak>, service_discovery: Box, + cancel: tokio_util::sync::CancellationToken, ) { - while let Some(cache) = cache.upgrade() { - tokio::time::sleep(service_discovery.interval()).await; + loop { + tokio::select! { + _ = tokio::time::sleep(service_discovery.interval()) => {} + _ = cancel.cancelled() => break, + } + let Some(cache) = cache.upgrade() else { break }; match service_discovery.pull_instances().await { Ok(instances) => { if let Err(error) = cache.set_peers(instances).await { @@ -50,3 +55,179 @@ pub(crate) async fn run_service_discovery( } } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::groupcache::ValueLoader; + use crate::options::Options; + use std::sync::atomic::{AtomicU32, Ordering}; + use tokio_util::sync::CancellationToken; + use std::sync::Arc; + use std::time::Duration; + + struct DummyLoader; + + #[async_trait] + impl ValueLoader for DummyLoader { + type Value = String; + async fn load( + &self, + key: &str, + ) -> Result> { + Ok(key.to_string()) + } + } + + #[test] + fn default_interval_is_ten_seconds() { + struct NoOp; + #[async_trait] + impl ServiceDiscovery for NoOp { + async fn pull_instances( + &self, + ) -> Result, Box> + { + Ok(HashSet::new()) + } + } + assert_eq!(NoOp.interval(), Duration::from_secs(10)); + } + + #[tokio::test] + async fn service_discovery_stops_when_cache_dropped() { + static PULL_COUNT: AtomicU32 = AtomicU32::new(0); + + struct CountingDiscovery; + #[async_trait] + impl ServiceDiscovery for CountingDiscovery { + async fn pull_instances( + &self, + ) -> Result, Box> + { + PULL_COUNT.fetch_add(1, Ordering::SeqCst); + Ok(HashSet::new()) + } + fn interval(&self) -> Duration { + Duration::from_millis(10) + } + } + + PULL_COUNT.store(0, Ordering::SeqCst); + + let addr: std::net::SocketAddr = "127.0.0.1:0".parse().unwrap(); + let peer = GroupcachePeer::from_socket(addr); + let inner = Arc::new(GroupcacheInner::new( + peer, + Box::new(DummyLoader), + CancellationToken::new(), + Options::default(), + )); + let weak = Arc::downgrade(&inner); + + let handle = tokio::spawn(run_service_discovery(weak, Box::new(CountingDiscovery), CancellationToken::new())); + + // Let it run a few iterations + tokio::time::sleep(Duration::from_millis(80)).await; + let count_before_drop = PULL_COUNT.load(Ordering::SeqCst); + assert!(count_before_drop > 0, "service discovery should have run"); + + // Drop the strong reference — next iteration should break + drop(inner); + tokio::time::sleep(Duration::from_millis(40)).await; + + assert!( + handle.is_finished(), + "task should stop after cache is dropped" + ); + } + + #[tokio::test] + async fn service_discovery_logs_set_peers_error() { + // Bind+drop so port is free but nothing listens. + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let unreachable_addr = listener.local_addr().unwrap(); + drop(listener); + + struct UnreachablePeerDiscovery(std::net::SocketAddr); + #[async_trait] + impl ServiceDiscovery for UnreachablePeerDiscovery { + async fn pull_instances( + &self, + ) -> Result, Box> + { + Ok(HashSet::from([GroupcachePeer::from_socket(self.0)])) + } + fn interval(&self) -> Duration { + Duration::from_millis(50) + } + } + + let addr: std::net::SocketAddr = "127.0.0.1:0".parse().unwrap(); + let peer = GroupcachePeer::from_socket(addr); + let inner = Arc::new(GroupcacheInner::new( + peer, + Box::new(DummyLoader), + CancellationToken::new(), + Options::default(), + )); + let weak = Arc::downgrade(&inner); + + let handle = tokio::spawn(run_service_discovery( + weak, + Box::new(UnreachablePeerDiscovery(unreachable_addr)), + CancellationToken::new(), + )); + + // Wait for at least one failed set_peers attempt. + tokio::time::sleep(Duration::from_millis(300)).await; + assert!( + !handle.is_finished(), + "task should keep running despite set_peers errors" + ); + + drop(inner); + tokio::time::sleep(Duration::from_millis(100)).await; + assert!(handle.is_finished()); + } + + #[tokio::test] + async fn service_discovery_logs_pull_error() { + struct FailingDiscovery; + #[async_trait] + impl ServiceDiscovery for FailingDiscovery { + async fn pull_instances( + &self, + ) -> Result, Box> + { + Err("discovery unavailable".into()) + } + fn interval(&self) -> Duration { + Duration::from_millis(10) + } + } + + let addr: std::net::SocketAddr = "127.0.0.1:0".parse().unwrap(); + let peer = GroupcachePeer::from_socket(addr); + let inner = Arc::new(GroupcacheInner::new( + peer, + Box::new(DummyLoader), + CancellationToken::new(), + Options::default(), + )); + let weak = Arc::downgrade(&inner); + + let handle = tokio::spawn(run_service_discovery(weak, Box::new(FailingDiscovery), CancellationToken::new())); + + // Let it attempt a few pulls (they all fail but shouldn't crash) + tokio::time::sleep(Duration::from_millis(50)).await; + assert!( + !handle.is_finished(), + "task should keep running despite errors" + ); + + drop(inner); + tokio::time::sleep(Duration::from_millis(30)).await; + assert!(handle.is_finished()); + } +} diff --git a/groupcache/src/status.rs b/groupcache/src/status.rs new file mode 100644 index 0000000..429349a --- /dev/null +++ b/groupcache/src/status.rs @@ -0,0 +1,47 @@ +//! Health and status reporting for groupcache. + +/// Health state of the groupcache instance. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum HealthState { + /// All systems operational, at least one peer connected. + Healthy, + /// Running but no peers connected yet. + /// Normal during startup or for single-node deployments. + NoPeers, + /// Shutting down — cancellation token was cancelled. + ShuttingDown, +} + +/// Snapshot of the groupcache instance's current state. +/// +/// Returned by [`Groupcache::status()`](crate::Groupcache::status). +/// The library reports state; the application decides policy. +/// +/// # Example +/// +/// ```no_run +/// # async fn example(groupcache: groupcache::Groupcache) { +/// use groupcache::status::HealthState; +/// +/// let status = groupcache.status(); +/// match status.health { +/// HealthState::Healthy | HealthState::NoPeers => { +/// // respond 200 OK to health check +/// } +/// HealthState::ShuttingDown => { +/// // respond 503 Service Unavailable +/// } +/// } +/// # } +/// ``` +#[derive(Debug, Clone)] +pub struct Status { + /// Overall health state. + pub health: HealthState, + /// Number of connected remote peers (excludes self). + pub peer_count: usize, + /// Approximate number of entries in the main cache (keys this node owns). + pub main_cache_size: u64, + /// Approximate number of entries in the hot cache (replicated keys). + pub hot_cache_size: u64, +} diff --git a/groupcache/tests/common/mod.rs b/groupcache/tests/common/mod.rs index 4db67ae..5da47d2 100644 --- a/groupcache/tests/common/mod.rs +++ b/groupcache/tests/common/mod.rs @@ -3,10 +3,12 @@ use anyhow::anyhow; use anyhow::Result; use async_trait::async_trait; -use groupcache::Groupcache; +use groupcache::{Groupcache, GroupcachePeer, ServiceDiscovery}; use moka::future::CacheBuilder; use pretty_assertions::assert_eq; use std::collections::HashMap; +use std::collections::HashSet; +use std::error::Error; use std::future::{pending, Future}; use std::sync::{Arc, RwLock}; use std::time::Duration; @@ -19,6 +21,11 @@ use tonic::transport::Server; pub static OS_ALLOCATED_PORT_ADDR: &str = "127.0.0.1:0"; pub static HOT_CACHE_TTL: Duration = Duration::from_millis(100); +/// Small delay to allow invalidation events to propagate over gRPC streams. +/// This is NOT a TTL wait -- streams deliver in sub-millisecond, but we need +/// to yield to the tokio runtime for the watcher task to process the event. +pub static INVALIDATION_PROPAGATION_DELAY: Duration = Duration::from_millis(50); + pub fn key_owned_by_instance(instance: TestGroupcache) -> String { format!("{}_0", instance.addr()) } @@ -102,6 +109,82 @@ pub async fn spawn_groupcache(instance_id: &str) -> Result { spawn_groupcache_instance(instance_id, OS_ALLOCATED_PORT_ADDR, pending()).await } +/// Spawn a groupcache instance with a ServiceDiscovery implementation attached. +pub async fn spawn_groupcache_with_service_discovery( + instance_id: &str, + sd: impl ServiceDiscovery + 'static, +) -> Result { + let listener = TcpListener::bind(OS_ALLOCATED_PORT_ADDR).await.unwrap(); + let addr = listener.local_addr()?; + let groupcache = Groupcache::builder(addr.into(), TestCacheLoader::new(instance_id), groupcache::CancellationToken::new()) + .hot_cache(CacheBuilder::default().time_to_live(HOT_CACHE_TTL).build()) + .enable_invalidation_streaming() + .service_discovery(sd) + .build(); + + let server = groupcache.grpc_service(); + tokio::spawn(async move { + Server::builder() + .add_service(server) + .serve_with_incoming_shutdown(TcpListenerStream::new(listener), pending::<()>()) + .await + .unwrap(); + }); + + Ok(groupcache) +} + +/// A test implementation of ServiceDiscovery that returns a configurable set of peers. +pub struct TestServiceDiscovery { + pub peers: Arc>>, + pub poll_interval: Duration, + pub error_on_next: Arc>, +} + +impl TestServiceDiscovery { + pub fn new(peers: HashSet, poll_interval: Duration) -> Self { + Self { + peers: Arc::new(RwLock::new(peers)), + poll_interval, + error_on_next: Arc::new(RwLock::new(false)), + } + } +} + +#[async_trait] +impl ServiceDiscovery for TestServiceDiscovery { + async fn pull_instances( + &self, + ) -> std::result::Result, Box> { + let should_error = { *self.error_on_next.read().unwrap() }; + if should_error { + return Err("Simulated service discovery error".into()); + } + let peers = self.peers.read().unwrap().clone(); + Ok(peers) + } + + fn interval(&self) -> Duration { + self.poll_interval + } +} + +/// A ServiceDiscovery that uses the default interval() implementation (10s). +/// Used to cover the default trait method. +pub struct DefaultIntervalServiceDiscovery { + pub peers: HashSet, +} + +#[async_trait] +impl ServiceDiscovery for DefaultIntervalServiceDiscovery { + async fn pull_instances( + &self, + ) -> std::result::Result, Box> { + Ok(self.peers.clone()) + } + // interval() is NOT overridden, exercising the default 10s implementation +} + pub async fn spawn_groupcache_instance( instance_id: &str, addr: &str, @@ -109,8 +192,9 @@ pub async fn spawn_groupcache_instance( ) -> Result { let listener = TcpListener::bind(addr).await.unwrap(); let addr = listener.local_addr()?; - let groupcache = Groupcache::builder(addr.into(), TestCacheLoader::new(instance_id)) + let groupcache = Groupcache::builder(addr.into(), TestCacheLoader::new(instance_id), groupcache::CancellationToken::new()) .hot_cache(CacheBuilder::default().time_to_live(HOT_CACHE_TTL).build()) + .enable_invalidation_streaming() .build(); let server = groupcache.grpc_service(); diff --git a/groupcache/tests/groupcache_test.rs b/groupcache/tests/groupcache_test.rs index 270c450..afc4ed8 100644 --- a/groupcache/tests/groupcache_test.rs +++ b/groupcache/tests/groupcache_test.rs @@ -45,7 +45,7 @@ async fn builder_api_test() { let loader = DummyLoader {}; let me: SocketAddr = "127.0.0.1:8080".parse().unwrap(); - let _groupcache = Groupcache::builder(me.into(), loader) + let _groupcache = Groupcache::builder(me.into(), loader, groupcache::CancellationToken::new()) .main_cache(main_cache) .hot_cache(hot_cache) .grpc_endpoint_builder(Box::new(|endpoint: Endpoint| { @@ -55,3 +55,20 @@ async fn builder_api_test() { .service_discovery(ServiceDiscoveryStub { me: me.into() }) .build(); } + +#[tokio::test] +async fn https_connection_attempt_uses_https_scheme() { + let me: SocketAddr = "127.0.0.1:0".parse().unwrap(); + let groupcache = Groupcache::builder(me.into(), DummyLoader {}, groupcache::CancellationToken::new()) + .https() + .build(); + + // Try to add a peer over HTTPS. The connection will fail (peer isn't TLS) + // but the code path that formats the "https://..." address is exercised. + let peer_addr: SocketAddr = "127.0.0.1:19876".parse().unwrap(); + let result = groupcache.add_peer(peer_addr.into()).await; + assert!( + result.is_err(), + "expected connection error for HTTPS to non-TLS peer" + ); +} diff --git a/groupcache/tests/integration_test.rs b/groupcache/tests/integration_test.rs index 87147d5..22ae5f4 100644 --- a/groupcache/tests/integration_test.rs +++ b/groupcache/tests/integration_test.rs @@ -8,7 +8,7 @@ use std::collections::HashSet; use std::net::SocketAddr; use std::ops::Sub; -use std::str::FromStr; +use std::time::Duration; use tokio::time; #[tokio::test] @@ -16,7 +16,7 @@ async fn test_when_there_is_only_one_peer_it_should_handle_entire_key_space() -> let groupcache = { let loader = TestCacheLoader::new("1"); let addr: SocketAddr = "127.0.0.1:8080".parse()?; - Groupcache::::builder(addr.into(), loader).build() + Groupcache::::builder(addr.into(), loader, groupcache::CancellationToken::new()).build() }; let key = "K-some-random-key-d2k"; @@ -75,7 +75,13 @@ async fn test_healthy_peers_added_despite_one_unhealthy() -> Result<()> { .map(|e| e.addr().into()) .collect::>(); - peers.insert(SocketAddr::from_str("127.0.0.1:8080")?.into()); + // Bind a port, then immediately drop the listener so the port is free + // but nothing is listening. This avoids flakiness from hardcoded ports + // that may be occupied by other processes. + let unreachable_listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?; + let unreachable_addr = unreachable_listener.local_addr()?; + drop(unreachable_listener); + peers.insert(unreachable_addr.into()); let result = me.set_peers(peers.clone()).await; if result.is_ok() { panic!("Expected failure to set peers with unreachable address"); @@ -236,12 +242,17 @@ async fn when_kv_is_loaded_it_should_be_cached_in_hot_cache() -> Result<()> { instance_two.remove(&key_on_instance_2).await?; + // With streaming invalidation, hot cache is cleared near-realtime. + time::sleep(INVALIDATION_PROPAGATION_DELAY).await; + + // Value should be reloaded (load count = 2) because streaming + // invalidation cleared instance_one's hot cache. successful_get_opts( &key_on_instance_2, instance_one.clone(), GetAssertions { expected_instance_id: Some("2".into()), - expected_load_count: Some(1), + expected_load_count: Some(2), ..GetAssertions::default() }, ) @@ -314,3 +325,463 @@ async fn when_key_is_removed_then_it_should_be_removed_from_owner() -> Result<() Ok(()) } + +#[tokio::test] +async fn when_key_is_removed_hot_cache_on_other_peers_should_be_invalidated_via_stream( +) -> Result<()> { + let (instance_one, instance_two) = two_connected_instances().await?; + let key_on_instance_2 = key_owned_by_instance(instance_two.clone()); + + // Instance 1 fetches key from instance 2 -> cached in instance 1's hot cache + successful_get_opts( + &key_on_instance_2, + instance_one.clone(), + GetAssertions { + expected_instance_id: Some("2".into()), + expected_load_count: Some(1), + ..GetAssertions::default() + }, + ) + .await; + + // Instance 2 (owner) removes the key + instance_two.remove(&key_on_instance_2).await?; + + // Wait for invalidation event to propagate via stream + time::sleep(INVALIDATION_PROPAGATION_DELAY).await; + + // Instance 1 should now reload from instance 2 (hot cache was cleared by stream) + successful_get_opts( + &key_on_instance_2, + instance_one.clone(), + GetAssertions { + expected_instance_id: Some("2".into()), + expected_load_count: Some(2), + ..GetAssertions::default() + }, + ) + .await; + + Ok(()) +} + +#[tokio::test] +async fn when_non_owner_removes_key_all_hot_caches_should_be_invalidated() -> Result<()> { + let (instance_one, instance_two) = two_connected_instances().await?; + let key_on_instance_2 = key_owned_by_instance(instance_two.clone()); + + // Instance 1 fetches key -> cached in instance 1's hot cache + successful_get_opts( + &key_on_instance_2, + instance_one.clone(), + GetAssertions { + expected_instance_id: Some("2".into()), + expected_load_count: Some(1), + ..GetAssertions::default() + }, + ) + .await; + + // Instance 1 (non-owner) calls remove -> routes Remove RPC to instance 2 (owner) + // Instance 2 broadcasts invalidation event -> instance 1's watcher clears hot cache + instance_one.remove(&key_on_instance_2).await?; + + // Wait for broadcast propagation + time::sleep(INVALIDATION_PROPAGATION_DELAY).await; + + // Instance 1 re-fetches -> should reload since both its hot cache + // (cleared by remove() + stream) and instance 2's main cache (cleared by Remove RPC) are empty + successful_get_opts( + &key_on_instance_2, + instance_one.clone(), + GetAssertions { + expected_instance_id: Some("2".into()), + expected_load_count: Some(2), + ..GetAssertions::default() + }, + ) + .await; + + Ok(()) +} + +// --------------------------------------------------------------------------- +// service_discovery.rs: exercise run_service_discovery loop +// --------------------------------------------------------------------------- + +#[tokio::test] +async fn test_service_discovery_default_interval_is_used() -> Result<()> { + let sd = DefaultIntervalServiceDiscovery { + peers: HashSet::new(), + }; + + let instance = spawn_groupcache_with_service_discovery("sd-default", sd).await?; + + // Give it a moment to execute at least one pull and call interval(). + time::sleep(Duration::from_millis(50)).await; + + // Verify the instance is functional + let key = "sd-default-key"; + successful_get(key, Some("sd-default"), instance.clone()).await; + + Ok(()) +} + +#[tokio::test] +async fn test_service_discovery_with_unreachable_peer_logs_error() -> Result<()> { + let unreachable_peer: GroupcachePeer = "127.0.0.1:19999".parse::()?.into(); + let peers: HashSet = vec![unreachable_peer].into_iter().collect(); + let sd = TestServiceDiscovery::new(peers, Duration::from_millis(50)); + + let _instance = spawn_groupcache_with_service_discovery("sd-unreachable", sd).await?; + + time::sleep(Duration::from_millis(200)).await; + + Ok(()) +} + +#[tokio::test] +async fn test_service_discovery_pulls_peers_and_updates_routing() -> Result<()> { + let discovered_peer = spawn_groupcache("sd-peer").await?; + + let peers: HashSet = vec![discovered_peer.addr().into()].into_iter().collect(); + let sd = TestServiceDiscovery::new(peers, Duration::from_millis(50)); + + let instance = spawn_groupcache_with_service_discovery("sd-main", sd).await?; + + time::sleep(Duration::from_millis(200)).await; + + let key_on_discovered = key_owned_by_instance(discovered_peer.clone()); + successful_get(&key_on_discovered, Some("sd-peer"), instance.clone()).await; + + Ok(()) +} + +#[tokio::test] +async fn test_service_discovery_error_does_not_crash_loop() -> Result<()> { + let discovered_peer = spawn_groupcache("sd-err-peer").await?; + + let peers: HashSet = vec![discovered_peer.addr().into()].into_iter().collect(); + let sd = TestServiceDiscovery::new(peers.clone(), Duration::from_millis(50)); + let error_flag = sd.error_on_next.clone(); + + *error_flag.write().unwrap() = true; + + let instance = spawn_groupcache_with_service_discovery("sd-err-main", sd).await?; + + time::sleep(Duration::from_millis(200)).await; + + *error_flag.write().unwrap() = false; + time::sleep(Duration::from_millis(200)).await; + + let key_on_discovered = key_owned_by_instance(discovered_peer.clone()); + successful_get(&key_on_discovered, Some("sd-err-peer"), instance.clone()).await; + + Ok(()) +} + +// --------------------------------------------------------------------------- +// groupcache_inner.rs: hot cache hit path +// --------------------------------------------------------------------------- + +#[tokio::test] +async fn test_hot_cache_hit_returns_cached_value_without_reload() -> Result<()> { + let (instance_one, instance_two) = two_connected_instances().await?; + let key_on_instance_2 = key_owned_by_instance(instance_two.clone()); + + successful_get_opts( + &key_on_instance_2, + instance_one.clone(), + GetAssertions { + expected_instance_id: Some("2".into()), + expected_load_count: Some(1), + ..GetAssertions::default() + }, + ) + .await; + + successful_get_opts( + &key_on_instance_2, + instance_one.clone(), + GetAssertions { + expected_instance_id: Some("2".into()), + expected_load_count: Some(1), + ..GetAssertions::default() + }, + ) + .await; + + Ok(()) +} + +// --------------------------------------------------------------------------- +// groupcache_inner.rs: add_peer on already-added peer (early return) +// --------------------------------------------------------------------------- + +#[tokio::test] +async fn test_add_peer_twice_is_idempotent() -> Result<()> { + let (instance_one, instance_two) = two_instances().await?; + + instance_one.add_peer(instance_two.addr().into()).await?; + instance_one.add_peer(instance_two.addr().into()).await?; + + let key_on_instance_2 = key_owned_by_instance(instance_two.clone()); + successful_get(&key_on_instance_2, Some("2"), instance_one.clone()).await; + + Ok(()) +} + +// --------------------------------------------------------------------------- +// groupcache_inner.rs: remove_peer on non-existent peer (early return) +// --------------------------------------------------------------------------- + +#[tokio::test] +async fn test_remove_peer_not_present_is_noop() -> Result<()> { + let instance = spawn_groupcache("rm-noop").await?; + + let fake_peer = GroupcachePeer::from_socket("127.0.0.1:19999".parse()?); + instance.remove_peer(fake_peer).await?; + + Ok(()) +} + +// --------------------------------------------------------------------------- +// groupcache_inner.rs: set_peers with new successful connections +// --------------------------------------------------------------------------- + +#[tokio::test] +async fn test_set_peers_adds_new_peers_via_update_routing_table() -> Result<()> { + let peers = spawn_instances(3).await?; + let (me, others) = peers.split_first().unwrap(); + + let all_peers: HashSet = others.iter().map(|e| e.addr().into()).collect(); + me.set_peers(all_peers).await?; + + for (i, peer) in others.iter().enumerate() { + let key = key_owned_by_instance(peer.clone()); + let peer_index = &*(i + 1).to_string(); + successful_get(&key, Some(peer_index), me.clone()).await; + } + + Ok(()) +} + +// --------------------------------------------------------------------------- +// invalidation.rs: watcher reconnect after peer restart +// --------------------------------------------------------------------------- + +#[tokio::test] +async fn test_invalidation_watcher_flushes_hot_cache_on_peer_disconnect() -> Result<()> { + let (shutdown_signal, shutdown_recv) = tokio::sync::oneshot::channel::<()>(); + let (shutdown_done_s, shutdown_done_r) = tokio::sync::oneshot::channel::<()>(); + + async fn shutdown_proxy( + shutdown_signal: tokio::sync::oneshot::Receiver<()>, + shutdown_done: tokio::sync::oneshot::Sender<()>, + ) { + shutdown_signal.await.unwrap(); + shutdown_done.send(()).unwrap(); + } + + let instance_one = spawn_groupcache("disc-1").await?; + let instance_two = spawn_groupcache_instance( + "disc-2", + OS_ALLOCATED_PORT_ADDR, + shutdown_proxy(shutdown_recv, shutdown_done_s), + ) + .await?; + + instance_one.add_peer(instance_two.addr().into()).await?; + instance_two.add_peer(instance_one.addr().into()).await?; + + time::sleep(Duration::from_millis(100)).await; + + let key_on_two = key_owned_by_instance(instance_two.clone()); + successful_get_opts( + &key_on_two, + instance_one.clone(), + GetAssertions { + expected_instance_id: Some("disc-2".into()), + expected_load_count: Some(1), + ..GetAssertions::default() + }, + ) + .await; + + shutdown_signal.send(()).unwrap(); + shutdown_done_r.await.unwrap(); + + time::sleep(Duration::from_millis(1000)).await; + + let result = instance_one.get(&key_on_two).await; + match result { + Ok(v) => { + assert!( + v.contains("disc-1"), + "expected local fallback after disconnect, got: '{}'", + v + ); + } + Err(e) => { + assert!( + e.to_string().contains("Transport") || e.to_string().contains("Loading error"), + "unexpected error: '{}'", + e + ); + } + } + + time::sleep(Duration::from_millis(500)).await; + + Ok(()) +} + +// --------------------------------------------------------------------------- +// set_peers with no changes returns early +// --------------------------------------------------------------------------- + +#[tokio::test] +async fn test_set_peers_with_same_peers_is_noop() -> Result<()> { + let (instance_one, instance_two) = two_connected_instances().await?; + + let peers: HashSet = vec![instance_two.addr().into()].into_iter().collect(); + + instance_one.set_peers(peers).await?; + + let key_on_instance_2 = key_owned_by_instance(instance_two.clone()); + successful_get(&key_on_instance_2, Some("2"), instance_one.clone()).await; + + Ok(()) +} + +// --------------------------------------------------------------------------- +// groupcache_inner.rs: evict_main_cache_keys_not_owned after peer joins +// --------------------------------------------------------------------------- + +#[tokio::test] +async fn test_main_cache_evicts_keys_that_move_to_new_peer() -> Result<()> { + let instance_one = spawn_groupcache("evict-1").await?; + let instance_two = spawn_groupcache("evict-2").await?; + + for i in 0..50 { + let key = format!("K-evict-safe-{}", i); + let _ = instance_one.get(&key).await?; + } + + instance_one.add_peer(instance_two.addr().into()).await?; + + let key_on_two = key_owned_by_instance(instance_two.clone()); + successful_get(&key_on_two, Some("evict-2"), instance_one.clone()).await; + + Ok(()) +} + +#[tokio::test] +async fn test_readding_peer_via_set_peers_replaces_invalidation_watcher() -> Result<()> { + let peers = spawn_instances(3).await?; + let (me, others) = peers.split_first().unwrap(); + + let all_peers: HashSet = others.iter().map(|e| e.addr().into()).collect(); + me.set_peers(all_peers.clone()).await?; + + let subset: HashSet = vec![others[0].addr().into()].into_iter().collect(); + me.set_peers(subset).await?; + me.set_peers(all_peers).await?; + + time::sleep(Duration::from_millis(100)).await; + + for (i, peer) in others.iter().enumerate() { + let key = key_owned_by_instance(peer.clone()); + let peer_index = &*(i + 1).to_string(); + successful_get(&key, Some(peer_index), me.clone()).await; + } + + Ok(()) +} + +#[tokio::test] +async fn test_add_peer_is_idempotent() -> Result<()> { + let (instance_one, instance_two) = two_connected_instances().await?; + + // Adding the same peer again should be a no-op (not an error). + instance_one + .add_peer(GroupcachePeer::from_socket(instance_two.addr())) + .await?; + + // Routing should still work correctly. + let key = key_owned_by_instance(instance_two.clone()); + successful_get(&key, Some("2"), instance_one.clone()).await; + + Ok(()) +} + +#[tokio::test] +async fn test_set_peers_connects_to_genuinely_new_peers() -> Result<()> { + // Create instances WITHOUT connecting them through spawn_instances. + let instance_one = spawn_groupcache("1").await?; + let instance_two = spawn_groupcache("2").await?; + let instance_three = spawn_groupcache("3").await?; + + // Use set_peers to introduce both new peers at once. + let peers: HashSet = [instance_two.addr(), instance_three.addr()] + .into_iter() + .map(GroupcachePeer::from_socket) + .collect(); + + instance_one.set_peers(peers).await?; + + // Verify routing works through the newly connected peers. + let key_on_two = key_owned_by_instance(instance_two.clone()); + successful_get(&key_on_two, Some("2"), instance_one.clone()).await; + + let key_on_three = key_owned_by_instance(instance_three.clone()); + successful_get(&key_on_three, Some("3"), instance_one.clone()).await; + + Ok(()) +} + +#[tokio::test] +async fn test_remove_from_disconnected_peer_returns_transport_error() -> Result<()> { + let (instance_one, instance_two) = two_instances_with_one_disconnected().await?; + + // Find a key owned by the disconnected instance_two. + let key = key_owned_by_instance(instance_two.clone()); + + // The graceful shutdown keeps existing HTTP/2 connections alive briefly. + // Retry remove until the connection actually drops. + let mut last_err = None; + for _ in 0..20 { + match instance_one.remove(&key).await { + Err(e) => { + last_err = Some(e); + break; + } + Ok(()) => { + tokio::time::sleep(std::time::Duration::from_millis(50)).await; + } + } + } + + let err = last_err.expect("expected transport error for remove on disconnected peer"); + let err_string = err.to_string(); + assert!( + err_string.contains("Transport"), + "expected Transport error, got: '{}'", + err_string + ); + + Ok(()) +} + +#[tokio::test] +async fn test_remove_peer_that_is_not_known_is_noop() -> Result<()> { + let instance = single_instance().await?; + let unknown_addr: SocketAddr = "127.0.0.1:19999".parse()?; + + // Removing a peer that was never added should succeed silently. + instance + .remove_peer(GroupcachePeer::from_socket(unknown_addr)) + .await?; + + Ok(()) +} diff --git a/groupcache/tests/metrics_test.rs b/groupcache/tests/metrics_test.rs index 2a50b33..68e1516 100644 --- a/groupcache/tests/metrics_test.rs +++ b/groupcache/tests/metrics_test.rs @@ -192,10 +192,10 @@ impl MockRecorder { fn counter_value(&self, key: &str) -> Option { let counters = self.registered_counters.borrow(); - return counters.get(key).map(|c| { + counters.get(key).map(|c| { let guard = c.count.lock().unwrap(); *guard - }); + }) } } diff --git a/maelstrom-tests/.gitignore b/maelstrom-tests/.gitignore new file mode 100644 index 0000000..f352be1 --- /dev/null +++ b/maelstrom-tests/.gitignore @@ -0,0 +1,3 @@ +/target +/.maelstrom +store/ diff --git a/maelstrom-tests/Cargo.toml b/maelstrom-tests/Cargo.toml new file mode 100644 index 0000000..c6ca6fa --- /dev/null +++ b/maelstrom-tests/Cargo.toml @@ -0,0 +1,14 @@ +[package] +name = "groupcache-maelstrom" +version = "0.1.0" +edition = "2021" + +[dependencies] +maelstrom-node = "0.1.6" +tokio = { version = "1", features = ["full"] } +serde = { version = "1", features = ["derive"] } +serde_json = "1" +async-trait = "0.1" +hashring = "0.3" +log = "0.4" +env_logger = "0.11" diff --git a/maelstrom-tests/Makefile b/maelstrom-tests/Makefile new file mode 100644 index 0000000..c05961d --- /dev/null +++ b/maelstrom-tests/Makefile @@ -0,0 +1,78 @@ +MAELSTROM_VERSION := 0.2.3 +MAELSTROM_DIR := .maelstrom +MAELSTROM_URL := https://github.com/jepsen-io/maelstrom/releases/download/v$(MAELSTROM_VERSION)/maelstrom.tar.bz2 +MAELSTROM_SRC_URL := https://github.com/jepsen-io/maelstrom/archive/refs/tags/v$(MAELSTROM_VERSION).tar.gz +MAELSTROM := $(MAELSTROM_DIR)/maelstrom/maelstrom +CARGO_TARGET := $(or $(CARGO_TARGET_DIR),$(CURDIR)/target) +BIN := $(CARGO_TARGET)/release/groupcache-maelstrom + +.PHONY: build install-maelstrom test-maelstrom test-echo test-baseline test-partitions test-packet-loss test-combined serve-results clean + +build: + cargo build --release + +install-maelstrom: $(MAELSTROM) + +$(MAELSTROM): + @echo "Installing Maelstrom $(MAELSTROM_VERSION)..." + mkdir -p $(MAELSTROM_DIR) + curl -sL $(MAELSTROM_URL) | tar xj -C $(MAELSTROM_DIR) + @echo "Downloading Maelstrom source..." + curl -sL $(MAELSTROM_SRC_URL) | tar xz -C $(MAELSTROM_DIR) + @echo "Installing custom workload and checker..." + cp maelstrom/src/maelstrom/workload/cache_convergence.clj \ + $(MAELSTROM_DIR)/maelstrom-$(MAELSTROM_VERSION)/src/maelstrom/workload/ + mkdir -p $(MAELSTROM_DIR)/maelstrom-$(MAELSTROM_VERSION)/src/maelstrom/checker + cp maelstrom/src/maelstrom/checker/convergence.clj \ + $(MAELSTROM_DIR)/maelstrom-$(MAELSTROM_VERSION)/src/maelstrom/checker/ + @echo "Registering workload..." + bash scripts/patch-maelstrom.sh \ + $(MAELSTROM_DIR)/maelstrom-$(MAELSTROM_VERSION)/src/maelstrom/core.clj + @echo "Maelstrom installed and patched." + +test-echo: build install-maelstrom + $(MAELSTROM) test -w echo --bin $(BIN) --node-count 1 --time-limit 5 + +# Delta-t thresholds per scenario (milliseconds): +# baseline: 100ms — near-realtime under normal conditions +# partitions: 35000ms — bounded by hot cache TTL (30s) during partition +# packet-loss: 1000ms — brief delays from dropped messages +# combined: 35000ms — partition-dominated worst case + +test-baseline: build install-maelstrom + cd $(MAELSTROM_DIR)/maelstrom-$(MAELSTROM_VERSION) && \ + DELTA_T_MS=100 lein run test -w cache-convergence \ + --bin $(BIN) \ + --node-count 5 \ + --time-limit 30 \ + --rate 50 + +test-partitions: build install-maelstrom + cd $(MAELSTROM_DIR)/maelstrom-$(MAELSTROM_VERSION) && \ + DELTA_T_MS=35000 lein run test -w cache-convergence \ + --bin $(BIN) \ + --node-count 5 \ + --time-limit 60 \ + --rate 50 \ + --nemesis partition \ + --nemesis-interval 10 + +test-partitions-long: build install-maelstrom + cd $(MAELSTROM_DIR)/maelstrom-$(MAELSTROM_VERSION) && \ + DELTA_T_MS=35000 lein run test -w cache-convergence \ + --bin $(BIN) \ + --node-count 5 \ + --time-limit 120 \ + --rate 50 \ + --nemesis partition \ + --nemesis-interval 5 + +test-maelstrom: test-baseline test-partitions test-partitions-long + @echo "All Maelstrom tests complete." + +serve-results: install-maelstrom + cd $(MAELSTROM_DIR)/maelstrom-$(MAELSTROM_VERSION) && lein run serve + +clean: + cargo clean + rm -rf $(MAELSTROM_DIR) diff --git a/maelstrom-tests/README.md b/maelstrom-tests/README.md new file mode 100644 index 0000000..ce5dfb4 --- /dev/null +++ b/maelstrom-tests/README.md @@ -0,0 +1,109 @@ +# groupcache-maelstrom + +Maelstrom-based convergence testing for groupcache's streaming invalidation protocol. + +## What This Tests + +After a key is invalidated, all nodes in the cluster must stop serving stale data +within a bounded time window — even under network partitions and packet loss. + +## Prerequisites + +- Rust toolchain (`cargo`) +- Java 11+ (`java -version`) +- Leiningen (`brew install leiningen` or see https://leiningen.org/) +- Graphviz (`brew install graphviz` — for result visualization) +- Gnuplot (`brew install gnuplot` — for latency graphs) + +## Quick Start + +```bash +# Run all convergence tests +make test-maelstrom + +# View results in browser +make serve-results +``` + +## Test Scenarios + +| Command | Scenario | What it tests | +|---------|----------|--------------| +| `make test-baseline` | No faults, 5 nodes, 30s | Normal convergence latency | +| `make test-partitions` | Network partitions every 10s | Recovery after partition heals | +| `make test-packet-loss` | Random packet drops | Convergence under lossy network | +| `make test-combined` | Partitions + packet loss | Worst case behavior | + +## How It Works + +Each test node implements groupcache's core algorithm: +- **Consistent hashing** (40 vnodes) for key ownership +- **Main cache** (owner) + **hot cache** (replicas, 30s TTL) +- **Invalidation broadcast** from owner to all peers + +Communication between nodes goes through Maelstrom's simulated network, +which can inject partitions and packet loss. + +### Operations + +| Operation | Description | +|-----------|-------------| +| `load` | Force a key to be loaded — owner generates a versioned value | +| `read` | Read a key from any cache tier (main, hot, or fetch from owner) | +| `invalidate` | Remove a key — owner clears main cache and broadcasts to peers | + +### Convergence Checker + +A custom Maelstrom checker analyzes the complete operation history post-hoc: + +1. For each `invalidate(key)` at time T, finds all subsequent reads that returned the stale (pre-invalidation) value +2. Computes **convergence latency** = time of last stale read - T +3. Reports latency percentiles (p50, p95, p99, max) +4. **PASS** if no stale reads persist beyond the delta-t threshold (5s default) +5. **FAIL** if any stale read occurs after the threshold + +### Example Output + +``` +:valid? true +:convergence {:p50-ms 2.1, :p95-ms 12.3, :p99-ms 48.7, :max-ms 312.0, :count 1523} +:stale-reads {:total 47, :violations 0} +:delta-t-ms 5000.0 +``` + +## What This Proves + +- The invalidation protocol converges under network faults +- Measured convergence latency under various fault conditions +- No permanent stale data after faults heal + +## What This Does NOT Prove + +- Performance of the real gRPC transport (covered by groupcache integration tests) +- Linearizability (not claimed — groupcache is eventually consistent) +- Byzantine fault tolerance (not claimed) + +## Project Structure + +``` +groupcache-maelstrom/ +├── Cargo.toml # Rust project +├── Makefile # All test commands +├── src/ +│ ├── main.rs # Entry point +│ ├── lib.rs # Module exports +│ ├── protocol.rs # Message type definitions +│ ├── cache.rs # Consistent hashing + cache with TTL +│ ├── node.rs # Node state management +│ └── handler.rs # Message handlers +├── tests/ +│ └── cache_test.rs # Unit tests +├── maelstrom/ +│ └── src/maelstrom/ +│ ├── workload/ +│ │ └── cache_convergence.clj # Client + operation generator +│ └── checker/ +│ └── convergence.clj # Convergence latency analysis +└── scripts/ + └── patch-maelstrom.sh # Registers workload in Maelstrom +``` diff --git a/maelstrom-tests/maelstrom/src/maelstrom/checker/convergence.clj b/maelstrom-tests/maelstrom/src/maelstrom/checker/convergence.clj new file mode 100644 index 0000000..689ba56 --- /dev/null +++ b/maelstrom-tests/maelstrom/src/maelstrom/checker/convergence.clj @@ -0,0 +1,122 @@ +(ns maelstrom.checker.convergence + "Checks that cache invalidation converges within a bounded time. + + Algorithm: + 1. Track the 'current value' for each key (set by load, cleared by invalidate) + 2. After each invalidate, any read returning the old value is 'stale' + 3. Convergence latency = time of last stale read - time of invalidation + 4. Failure = stale read persists beyond delta-t after invalidation" + (:require [jepsen [checker :as checker] + [util :as util]])) + +(def default-delta-t-nanos + "Default convergence threshold in nanoseconds. + Override via DELTA_T_MS environment variable (milliseconds)." + (if-let [env-ms (System/getenv "DELTA_T_MS")] + (* (Long/parseLong env-ms) 1000 1000) + (* 100 1000 1000))) + +(defn completed-ops + "Filter to only completed (ok) operations." + [history] + (filter #(= :ok (:type %)) history)) + +(defn build-key-timeline + "Build a timeline of events per key from the history. + Returns {key [{:time t :op :load/:read/:invalidate :value v} ...]}" + [history] + (reduce + (fn [acc op] + (let [key (get-in op [:value :key]) + result (get-in op [:value :result]) + entry {:time (:time op) + :op (:f op) + :value result}] + (update acc key (fnil conj []) entry))) + {} + (completed-ops history))) + +(defn analyze-key + "Analyze convergence for a single key's timeline. + Returns a map with :stale-reads and :convergence-latencies." + [events] + (loop [remaining events + current-value nil + pre-inv-value nil + inv-time nil + stale-reads [] + conv-latencies []] + (if (empty? remaining) + {:stale-reads stale-reads + :convergence-latencies conv-latencies} + (let [event (first remaining) + rest (rest remaining)] + (case (:op event) + :load + (recur rest (:value event) pre-inv-value inv-time stale-reads conv-latencies) + + :invalidate + (recur rest nil current-value (:time event) stale-reads conv-latencies) + + :read + (if (and inv-time + pre-inv-value + (:value event) + (= (:value event) pre-inv-value)) + ;; Stale read detected + (let [latency-ns (- (:time event) inv-time)] + (recur rest current-value pre-inv-value inv-time + (conj stale-reads {:time (:time event) + :latency-ns latency-ns + :value (:value event)}) + (conj conv-latencies latency-ns))) + ;; Normal read + (recur rest current-value pre-inv-value inv-time stale-reads conv-latencies)) + + ;; Unknown op, skip + (recur rest current-value pre-inv-value inv-time stale-reads conv-latencies)))))) + +(defn percentile + "Compute the p-th percentile of a sorted sequence." + [sorted-vals p] + (if (empty? sorted-vals) + 0 + (let [idx (min (dec (count sorted-vals)) + (int (Math/ceil (* (/ p 100.0) (count sorted-vals)))))] + (nth sorted-vals idx)))) + +(defn ns->ms + "Convert nanoseconds to milliseconds (double)." + [ns] + (/ (double ns) 1000000.0)) + +(defn checker + "Create a convergence checker. + Options: + :delta-t-nanos Maximum allowed convergence time (default 5s)" + ([] (checker {})) + ([opts] + (let [delta-t (get opts :delta-t-nanos default-delta-t-nanos)] + (reify checker/Checker + (check [this test history check-opts] + (let [timelines (build-key-timeline history) + analyses (map (fn [[k events]] [k (analyze-key events)]) timelines) + all-latencies (->> analyses + (mapcat (fn [[_ a]] (:convergence-latencies a))) + sort + vec) + all-stale (->> analyses + (mapcat (fn [[_ a]] (:stale-reads a))) + vec) + violations (filter #(> (:latency-ns %) delta-t) all-stale) + valid? (empty? violations)] + {:valid? valid? + :convergence (when (seq all-latencies) + {:p50-ms (ns->ms (percentile all-latencies 50)) + :p95-ms (ns->ms (percentile all-latencies 95)) + :p99-ms (ns->ms (percentile all-latencies 99)) + :max-ms (ns->ms (last all-latencies)) + :count (count all-latencies)}) + :stale-reads {:total (count all-stale) + :violations (count violations)} + :delta-t-ms (ns->ms delta-t)})))))) diff --git a/maelstrom-tests/maelstrom/src/maelstrom/workload/cache_convergence.clj b/maelstrom-tests/maelstrom/src/maelstrom/workload/cache_convergence.clj new file mode 100644 index 0000000..bdf6347 --- /dev/null +++ b/maelstrom-tests/maelstrom/src/maelstrom/workload/cache_convergence.clj @@ -0,0 +1,97 @@ +(ns maelstrom.workload.cache-convergence + "A workload for testing cache invalidation convergence. + + Operations: + - load: Force a key to be loaded (generates a new versioned value) + - read: Read a key from any cache tier + - invalidate: Remove a key from all caches + + The checker verifies that after invalidation + delta-t, all nodes + converge (no stale reads persist)." + (:require [maelstrom [client :as c] + [net :as net]] + [maelstrom.checker.convergence :as convergence] + [jepsen [checker :as checker] + [client :as client] + [generator :as gen]] + [schema.core :as s] + [slingshot.slingshot :refer [try+ throw+]])) + +(def key-count + "Number of distinct keys to test." + 10) + +;; ── RPC definitions ───────────────────────────────────────────────── + +(c/defrpc load-key! + "Force a key to be loaded on the owning node." + {:type (s/eq "load") + :key s/Str} + {:type (s/eq "load_ok") + :value s/Any}) + +(c/defrpc read-key + "Read a key from any cache tier." + {:type (s/eq "read") + :key s/Str} + {:type (s/eq "read_ok") + :value s/Any}) + +(c/defrpc invalidate-key! + "Invalidate a key across all nodes." + {:type (s/eq "invalidate") + :key s/Str} + {:type (s/eq "invalidate_ok")}) + +;; ── Client ────────────────────────────────────────────────────────── + +(defn client + "Construct a cache-convergence client for the given network." + ([net] + (client net nil nil)) + ([net conn node] + (reify client/Client + (open! [this test node] + (client net (c/open! net) node)) + + (setup! [this test]) + + (invoke! [_ test op] + (let [k (get-in op [:value :key])] + (c/with-errors op #{:read} + (case (:f op) + :load + (let [res (load-key! conn node {:key k})] + (assoc op :type :ok + :value (assoc (:value op) :result (:value res)))) + + :read + (let [res (read-key conn node {:key k})] + (assoc op :type :ok + :value (assoc (:value op) :result (:value res)))) + + :invalidate + (do (invalidate-key! conn node {:key k}) + (assoc op :type :ok)))))) + + (teardown! [_ test]) + + (close! [_ test] + (c/close! conn))))) + +;; ── Generator ─────────────────────────────────────────────────────── + +(defn load-op [_ _] {:type :invoke, :f :load, :value {:key (str "k" (rand-int key-count))}}) +(defn read-op [_ _] {:type :invoke, :f :read, :value {:key (str "k" (rand-int key-count))}}) +(defn inv-op [_ _] {:type :invoke, :f :invalidate, :value {:key (str "k" (rand-int key-count))}}) + +;; ── Workload ──────────────────────────────────────────────────────── + +(defn workload + "Constructs a workload for cache convergence testing. + + {:net A Maelstrom network}" + [opts] + {:client (client (:net opts)) + :generator (gen/mix [load-op read-op read-op read-op inv-op]) + :checker (convergence/checker)}) diff --git a/maelstrom-tests/scripts/patch-maelstrom.sh b/maelstrom-tests/scripts/patch-maelstrom.sh new file mode 100755 index 0000000..1ff3d56 --- /dev/null +++ b/maelstrom-tests/scripts/patch-maelstrom.sh @@ -0,0 +1,39 @@ +#!/usr/bin/env bash +set -euo pipefail + +CORE_FILE="$1" + +if [ -z "$CORE_FILE" ] || [ ! -f "$CORE_FILE" ]; then + echo "Usage: $0 " + exit 1 +fi + +# Check if already patched +if grep -q "cache-convergence" "$CORE_FILE"; then + echo "Already patched." + exit 0 +fi + +# The require block uses nested form: +# [maelstrom.workload [broadcast :as broadcast] +# [echo :as echo] +# ... +# [unique-ids :as unique-ids]] +# +# Add our workload after unique-ids, keeping the closing ]] +sed -i.bak 's/\[unique-ids :as unique-ids\]\]/[unique-ids :as unique-ids]\ + [cache-convergence :as cache-convergence]]/' "$CORE_FILE" + +# The workloads map uses keywords: +# {:broadcast broadcast/workload +# ... +# :unique-ids unique-ids/workload}) +# +# Add our entry before the closing }) +sed -i.bak 's/:unique-ids unique-ids\/workload})/:unique-ids unique-ids\/workload\ + :cache-convergence cache-convergence\/workload})/' "$CORE_FILE" + +# Clean up backup files +rm -f "${CORE_FILE}.bak" + +echo "Patched $CORE_FILE with cache-convergence workload." diff --git a/maelstrom-tests/src/cache.rs b/maelstrom-tests/src/cache.rs new file mode 100644 index 0000000..6c7fea9 --- /dev/null +++ b/maelstrom-tests/src/cache.rs @@ -0,0 +1,86 @@ +use hashring::HashRing; +use std::collections::HashMap; +use std::time::{Duration, Instant}; +use tokio::sync::RwLock; + +const VNODES_PER_NODE: usize = 40; + +#[derive(Clone)] +pub struct ConsistentHashRing { + ring: HashRing, + nodes: Vec, +} + +impl ConsistentHashRing { + pub fn new(node_ids: &[String]) -> Self { + let mut ring = HashRing::new(); + for node_id in node_ids { + for i in 0..VNODES_PER_NODE { + ring.add(format!("{}_{}", node_id, i)); + } + } + Self { + ring, + nodes: node_ids.to_vec(), + } + } + + pub fn owner(&self, key: &str) -> &str { + let vnode = self.ring.get(&key).expect("ring cannot be empty"); + let node_id = vnode.split('_').next().unwrap(); + self.nodes.iter().find(|n| n.as_str() == node_id).unwrap() + } +} + +struct CacheEntry { + value: String, + inserted_at: Instant, +} + +pub struct Cache { + entries: RwLock>, + ttl: Option, +} + +impl Cache { + pub fn new(ttl: Option) -> Self { + Self { + entries: RwLock::new(HashMap::new()), + ttl, + } + } + + pub async fn get(&self, key: &str) -> Option { + let entries = self.entries.read().await; + if let Some(entry) = entries.get(key) { + if let Some(ttl) = self.ttl { + if entry.inserted_at.elapsed() > ttl { + drop(entries); + self.entries.write().await.remove(key); + return None; + } + } + Some(entry.value.clone()) + } else { + None + } + } + + pub async fn insert(&self, key: String, value: String) { + self.entries.write().await.insert( + key, + CacheEntry { + value, + inserted_at: Instant::now(), + }, + ); + } + + pub async fn remove(&self, key: &str) { + self.entries.write().await.remove(key); + } + + pub async fn invalidate_all(&self) { + self.entries.write().await.clear(); + } +} diff --git a/maelstrom-tests/src/handler.rs b/maelstrom-tests/src/handler.rs new file mode 100644 index 0000000..235ff5c --- /dev/null +++ b/maelstrom-tests/src/handler.rs @@ -0,0 +1,245 @@ +use crate::node::GroupcacheNode; +use crate::protocol::*; +use log::warn; +use maelstrom::protocol::{ErrorMessageBody, Message}; +use maelstrom::{done, Result, Runtime}; +use std::sync::Arc; + +/// Top-level message dispatcher. Routes incoming messages to the appropriate handler +/// based on message type. +pub async fn handle_message( + node: &Arc, + runtime: &Runtime, + req: &Message, +) -> Result<()> { + match req.get_type() { + "load" => handle_load(node, runtime, req).await, + "read" => handle_read(node, runtime, req).await, + "invalidate" => handle_invalidate(node, runtime, req).await, + "gc_get" => handle_gc_get(node, runtime, req).await, + "gc_load" => handle_gc_load(node, runtime, req).await, + "gc_remove" => handle_gc_remove(node, runtime, req).await, + "gc_invalidation_event" => handle_gc_invalidation_event(node, req).await, + _ => done(runtime.clone(), req.clone()), + } +} + +/// Client operation: load a value into the cache for a given key. +/// +/// - If this node is the owner: generate a new value, store in main_cache. +/// - If not owner: RPC `gc_load` to the owner, cache the result in hot_cache. +/// - Reply with `load_ok` + value. +async fn handle_load( + node: &Arc, + runtime: &Runtime, + req: &Message, +) -> Result<()> { + let load_req: LoadRequest = req.body.as_obj()?; + let key = load_req.key; + + let value = if node.is_owner(&key) { + // We are the owner: generate and store locally. + let val = node.next_value(&key).await; + node.main_cache.insert(key.clone(), val.clone()).await; + val + } else { + // Forward to owner via gc_load RPC. + let owner = node.owner_of(&key).to_string(); + let gc_load = GcLoad::new(key.clone()); + let rpc_result = runtime.rpc(owner, gc_load).await?; + let resp_msg: Message = rpc_result.await?; + let gc_load_ok: GcLoadOk = resp_msg.body.as_obj()?; + // Cache in hot_cache for future reads. + node.hot_cache + .insert(key.clone(), gc_load_ok.value.clone()) + .await; + gc_load_ok.value + }; + + runtime.reply(req.clone(), LoadResponse::new(value)).await +} + +/// Client operation: read a cached value for a given key. +/// +/// - Check local caches (main_cache then hot_cache). +/// - If hit: reply with value. +/// - If miss and we are the owner: reply with error code 20 (key not found). +/// - If miss and remote owner: RPC `gc_get` to owner, cache in hot_cache, reply. +/// - If RPC fails: reply with error code 11 (temporarily unavailable). +async fn handle_read( + node: &Arc, + runtime: &Runtime, + req: &Message, +) -> Result<()> { + let read_req: ReadRequest = req.body.as_obj()?; + let key = read_req.key; + + // Check local caches first. + if let Some(value) = node.local_lookup(&key).await { + return runtime + .reply(req.clone(), ReadResponse::new(value)) + .await; + } + + if node.is_owner(&key) { + // We are the owner but don't have the key -- not found. + return runtime + .reply( + req.clone(), + ErrorMessageBody::new(20, "key not found"), + ) + .await; + } + + // Forward to owner via gc_get RPC. + let owner = node.owner_of(&key).to_string(); + let gc_get = GcGet::new(key.clone()); + let rpc_result = runtime.rpc(owner, gc_get).await; + + match rpc_result { + Ok(call) => match call.await { + Ok(resp_msg) => { + let gc_get_ok: GcGetOk = resp_msg.body.as_obj()?; + match gc_get_ok.value { + Some(value) => { + node.hot_cache.insert(key.clone(), value.clone()).await; + runtime + .reply(req.clone(), ReadResponse::new(value)) + .await + } + None => { + runtime + .reply( + req.clone(), + ErrorMessageBody::new(20, "key not found"), + ) + .await + } + } + } + Err(e) => { + warn!("gc_get RPC failed: {}", e); + runtime + .reply( + req.clone(), + ErrorMessageBody::new(11, "temporarily unavailable"), + ) + .await + } + }, + Err(e) => { + warn!("gc_get RPC send failed: {}", e); + runtime + .reply( + req.clone(), + ErrorMessageBody::new(11, "temporarily unavailable"), + ) + .await + } + } +} + +/// Client operation: invalidate a cached entry for a given key. +/// +/// - Clear own hot_cache for the key. +/// - If owner: clear main_cache, broadcast `gc_invalidation_event` to all peers. +/// - If not owner: RPC `gc_remove` to the owner. +/// - Reply with `invalidate_ok`. +async fn handle_invalidate( + node: &Arc, + runtime: &Runtime, + req: &Message, +) -> Result<()> { + let inv_req: InvalidateRequest = req.body.as_obj()?; + let key = inv_req.key; + + // Always clear our own hot_cache. + node.hot_cache.remove(&key).await; + + if node.is_owner(&key) { + // Clear main_cache and broadcast invalidation to peers. + node.main_cache.remove(&key).await; + broadcast_invalidation(node, runtime, &key); + } else { + // Forward to owner via gc_remove RPC. + let owner = node.owner_of(&key).to_string(); + let gc_remove = GcRemove::new(key.clone()); + let rpc_result = runtime.rpc(owner, gc_remove).await?; + let _resp: Message = rpc_result.await?; + } + + runtime + .reply(req.clone(), InvalidateResponse::new()) + .await +} + +/// Inter-node: peer requests a value from our main_cache. +/// We are the owner for this key. +async fn handle_gc_get( + node: &Arc, + runtime: &Runtime, + req: &Message, +) -> Result<()> { + let gc_get: GcGet = req.body.as_obj()?; + let key = gc_get.key; + let value = node.main_cache.get(&key).await; + runtime + .reply(req.clone(), GcGetOk::new(key, value)) + .await +} + +/// Inter-node: peer asks us (the owner) to load/generate a value. +async fn handle_gc_load( + node: &Arc, + runtime: &Runtime, + req: &Message, +) -> Result<()> { + let gc_load: GcLoad = req.body.as_obj()?; + let key = gc_load.key; + let value = node.next_value(&key).await; + node.main_cache.insert(key.clone(), value.clone()).await; + runtime + .reply(req.clone(), GcLoadOk::new(key, value)) + .await +} + +/// Inter-node: peer asks us (the owner) to remove a key. +/// Clear main_cache and broadcast invalidation to all peers. +async fn handle_gc_remove( + node: &Arc, + runtime: &Runtime, + req: &Message, +) -> Result<()> { + let gc_remove: GcRemove = req.body.as_obj()?; + let key = gc_remove.key; + node.main_cache.remove(&key).await; + broadcast_invalidation(node, runtime, &key); + runtime + .reply(req.clone(), GcRemoveOk::new(key)) + .await +} + +/// Inter-node fire-and-forget: a peer tells us that a key has been invalidated. +/// Clear our hot_cache for that key. No reply needed. +async fn handle_gc_invalidation_event( + node: &Arc, + req: &Message, +) -> Result<()> { + let event: GcInvalidationEvent = req.body.as_obj()?; + node.hot_cache.remove(&event.key).await; + Ok(()) +} + +/// Broadcast a `gc_invalidation_event` to all peers (fire-and-forget). +fn broadcast_invalidation( + _node: &Arc, + runtime: &Runtime, + key: &str, +) { + for peer in runtime.neighbours() { + let event = GcInvalidationEvent::new(key.to_string()); + if let Err(e) = runtime.send_async(peer.as_str(), event) { + warn!("Failed to send gc_invalidation_event to {}: {}", peer, e); + } + } +} diff --git a/maelstrom-tests/src/lib.rs b/maelstrom-tests/src/lib.rs new file mode 100644 index 0000000..8dccf27 --- /dev/null +++ b/maelstrom-tests/src/lib.rs @@ -0,0 +1,4 @@ +pub mod cache; +pub mod handler; +pub mod node; +pub mod protocol; diff --git a/maelstrom-tests/src/main.rs b/maelstrom-tests/src/main.rs new file mode 100644 index 0000000..506223b --- /dev/null +++ b/maelstrom-tests/src/main.rs @@ -0,0 +1,37 @@ +use async_trait::async_trait; +use groupcache_maelstrom::handler::handle_message; +use groupcache_maelstrom::node::GroupcacheNode; +use maelstrom::protocol::Message; +use maelstrom::{Node, Result, Runtime}; +use std::sync::Arc; +use tokio::sync::OnceCell; + +static NODE: OnceCell> = OnceCell::const_new(); + +#[derive(Clone, Default)] +struct Handler; + +#[async_trait] +impl Node for Handler { + async fn process(&self, runtime: Runtime, req: Message) -> Result<()> { + let node = NODE + .get_or_init(|| async { + let node_id = runtime.node_id().to_string(); + let node_ids: Vec = + runtime.nodes().iter().map(|s| s.to_string()).collect(); + Arc::new(GroupcacheNode::new(node_id, node_ids)) + }) + .await; + + handle_message(node, &runtime, &req).await + } +} + +fn main() -> Result<()> { + Runtime::init(try_main()) +} + +async fn try_main() -> Result<()> { + let handler = Arc::new(Handler); + Runtime::new().with_handler(handler).run().await +} diff --git a/maelstrom-tests/src/node.rs b/maelstrom-tests/src/node.rs new file mode 100644 index 0000000..32b0ed4 --- /dev/null +++ b/maelstrom-tests/src/node.rs @@ -0,0 +1,51 @@ +use crate::cache::{Cache, ConsistentHashRing}; +use std::collections::HashMap; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::time::Duration; +use tokio::sync::RwLock; + +const HOT_CACHE_TTL: Duration = Duration::from_secs(30); + +pub struct GroupcacheNode { + pub node_id: String, + pub ring: ConsistentHashRing, + pub main_cache: Cache, + pub hot_cache: Cache, + versions: RwLock>, +} + +impl GroupcacheNode { + pub fn new(node_id: String, all_node_ids: Vec) -> Self { + Self { + ring: ConsistentHashRing::new(&all_node_ids), + node_id, + main_cache: Cache::new(None), + hot_cache: Cache::new(Some(HOT_CACHE_TTL)), + versions: RwLock::new(HashMap::new()), + } + } + + pub fn is_owner(&self, key: &str) -> bool { + self.ring.owner(key) == self.node_id + } + + pub fn owner_of(&self, key: &str) -> &str { + self.ring.owner(key) + } + + pub async fn next_value(&self, key: &str) -> String { + let mut versions = self.versions.write().await; + let counter = versions + .entry(key.to_string()) + .or_insert_with(|| AtomicU64::new(0)); + let version = counter.fetch_add(1, Ordering::SeqCst) + 1; + format!("{}:{}", self.node_id, version) + } + + pub async fn local_lookup(&self, key: &str) -> Option { + if let Some(v) = self.main_cache.get(key).await { + return Some(v); + } + self.hot_cache.get(key).await + } +} diff --git a/maelstrom-tests/src/protocol.rs b/maelstrom-tests/src/protocol.rs new file mode 100644 index 0000000..b0b61b6 --- /dev/null +++ b/maelstrom-tests/src/protocol.rs @@ -0,0 +1,205 @@ +use serde::{Deserialize, Serialize}; + +// --------------------------------------------------------------------------- +// Client operations (Maelstrom workload -> node) +// --------------------------------------------------------------------------- + +/// Request to load a value into the cache for a given key. +#[derive(Debug, Deserialize)] +pub struct LoadRequest { + pub key: String, +} + +/// Successful response to a `load` request. +#[derive(Debug, Serialize)] +pub struct LoadResponse { + #[serde(rename = "type")] + pub msg_type: String, + pub value: String, +} + +impl LoadResponse { + pub fn new(value: String) -> Self { + Self { + msg_type: "load_ok".to_string(), + value, + } + } +} + +/// Request to read a cached value for a given key. +#[derive(Debug, Deserialize)] +pub struct ReadRequest { + pub key: String, +} + +/// Successful response to a `read` request. +#[derive(Debug, Serialize)] +pub struct ReadResponse { + #[serde(rename = "type")] + pub msg_type: String, + pub value: String, +} + +impl ReadResponse { + pub fn new(value: String) -> Self { + Self { + msg_type: "read_ok".to_string(), + value, + } + } +} + +/// Request to invalidate a cached entry for a given key. +#[derive(Debug, Deserialize)] +pub struct InvalidateRequest { + pub key: String, +} + +/// Successful response to an `invalidate` request. +#[derive(Debug, Serialize)] +pub struct InvalidateResponse { + #[serde(rename = "type")] + pub msg_type: String, +} + +impl Default for InvalidateResponse { + fn default() -> Self { + Self { + msg_type: "invalidate_ok".to_string(), + } + } +} + +impl InvalidateResponse { + pub fn new() -> Self { + Self::default() + } +} + +// --------------------------------------------------------------------------- +// Inter-node messages (node <-> node, groupcache protocol) +// --------------------------------------------------------------------------- + +/// Request a value from a peer's local cache. +#[derive(Debug, Serialize, Deserialize)] +pub struct GcGet { + #[serde(rename = "type")] + pub msg_type: String, + pub key: String, +} + +impl GcGet { + pub fn new(key: String) -> Self { + Self { + msg_type: "gc_get".to_string(), + key, + } + } +} + +/// Response to `gc_get` — the value may be absent if the peer doesn't have it. +#[derive(Debug, Serialize, Deserialize)] +pub struct GcGetOk { + #[serde(rename = "type")] + pub msg_type: String, + pub key: String, + pub value: Option, +} + +impl GcGetOk { + pub fn new(key: String, value: Option) -> Self { + Self { + msg_type: "gc_get_ok".to_string(), + key, + value, + } + } +} + +/// Ask a peer to load (fetch from origin) a value for the given key. +#[derive(Debug, Serialize, Deserialize)] +pub struct GcLoad { + #[serde(rename = "type")] + pub msg_type: String, + pub key: String, +} + +impl GcLoad { + pub fn new(key: String) -> Self { + Self { + msg_type: "gc_load".to_string(), + key, + } + } +} + +/// Response to `gc_load` with the loaded value. +#[derive(Debug, Serialize, Deserialize)] +pub struct GcLoadOk { + #[serde(rename = "type")] + pub msg_type: String, + pub key: String, + pub value: String, +} + +impl GcLoadOk { + pub fn new(key: String, value: String) -> Self { + Self { + msg_type: "gc_load_ok".to_string(), + key, + value, + } + } +} + +/// Ask a peer to remove a key from its local cache. +#[derive(Debug, Serialize, Deserialize)] +pub struct GcRemove { + #[serde(rename = "type")] + pub msg_type: String, + pub key: String, +} + +impl GcRemove { + pub fn new(key: String) -> Self { + Self { + msg_type: "gc_remove".to_string(), + key, + } + } +} + +/// Acknowledgement that the key was removed. +#[derive(Debug, Serialize, Deserialize)] +pub struct GcRemoveOk { + #[serde(rename = "type")] + pub msg_type: String, + pub key: String, +} + +impl GcRemoveOk { + pub fn new(key: String) -> Self { + Self { + msg_type: "gc_remove_ok".to_string(), + key, + } + } +} + +/// Broadcast event telling peers that a key has been invalidated. +#[derive(Debug, Serialize, Deserialize)] +pub struct GcInvalidationEvent { + #[serde(rename = "type")] + pub msg_type: String, + pub key: String, +} + +impl GcInvalidationEvent { + pub fn new(key: String) -> Self { + Self { + msg_type: "gc_invalidation_event".to_string(), + key, + } + } +} diff --git a/maelstrom-tests/tests/cache_test.rs b/maelstrom-tests/tests/cache_test.rs new file mode 100644 index 0000000..1bb84ea --- /dev/null +++ b/maelstrom-tests/tests/cache_test.rs @@ -0,0 +1,64 @@ +use groupcache_maelstrom::cache::{Cache, ConsistentHashRing}; +use std::time::Duration; + +#[test] +fn hash_ring_deterministic_ownership() { + let nodes = vec!["n1".into(), "n2".into(), "n3".into()]; + let ring = ConsistentHashRing::new(&nodes); + let owner1 = ring.owner("key-abc"); + let owner2 = ring.owner("key-abc"); + assert_eq!(owner1, owner2); +} + +#[test] +fn hash_ring_distributes_keys() { + let nodes = vec!["n1".into(), "n2".into(), "n3".into()]; + let ring = ConsistentHashRing::new(&nodes); + let mut owners = std::collections::HashSet::new(); + for i in 0..100 { + owners.insert(ring.owner(&format!("key-{}", i)).to_string()); + } + assert_eq!(owners.len(), 3); +} + +#[test] +fn hash_ring_owner_is_valid_node() { + let nodes = vec!["n1".into(), "n2".into(), "n3".into()]; + let ring = ConsistentHashRing::new(&nodes); + let owner = ring.owner("any-key"); + assert!(nodes.contains(&owner.to_string())); +} + +#[tokio::test] +async fn cache_insert_and_get() { + let cache = Cache::new(None); + cache.insert("k1".into(), "v1".into()).await; + assert_eq!(cache.get("k1").await, Some("v1".into())); +} + +#[tokio::test] +async fn cache_remove() { + let cache = Cache::new(None); + cache.insert("k1".into(), "v1".into()).await; + cache.remove("k1").await; + assert_eq!(cache.get("k1").await, None); +} + +#[tokio::test] +async fn cache_ttl_expiry() { + let cache = Cache::new(Some(Duration::from_millis(50))); + cache.insert("k1".into(), "v1".into()).await; + assert_eq!(cache.get("k1").await, Some("v1".into())); + tokio::time::sleep(Duration::from_millis(100)).await; + assert_eq!(cache.get("k1").await, None); +} + +#[tokio::test] +async fn cache_invalidate_all() { + let cache = Cache::new(None); + cache.insert("k1".into(), "v1".into()).await; + cache.insert("k2".into(), "v2".into()).await; + cache.invalidate_all().await; + assert_eq!(cache.get("k1").await, None); + assert_eq!(cache.get("k2").await, None); +}