alternate availability store schema (#2237)

* alternate availability store schema

* improvements

* tweaks

* new DB schema and skeleton

* expand skeleton and tweaks

* handle backing and inclusion

* let finality be handled later

* handle finalized blocks

* implement query methods

* implement chunk storing

* StoreAvailableData

* fix an off-by-one

* implement pruning

* reinstate subsystem trait impl

* reinstate metrics

* fix warnings

* remove chunks_cache

* oops

* actually store the available data

* mockable pruning interval

* fix tests

* spacing

* fix code grumbles

* guide improvements

* make time mockable

* implement a mocked clock for testing

* return DB errors

* Update node/core/av-store/Cargo.toml

Co-authored-by: Bastian Köcher <bkchr@users.noreply.github.com>

* Update roadmap/implementers-guide/src/node/utility/availability-store.md

Co-authored-by: Bastian Köcher <bkchr@users.noreply.github.com>

* Update roadmap/implementers-guide/src/node/utility/availability-store.md

Co-authored-by: Bastian Köcher <bkchr@users.noreply.github.com>

* review grumbles & clarity

* fix review grumbles

* Add docs

Co-authored-by: Andronik Ordian <write@reusable.software>

Co-authored-by: Bastian Köcher <bkchr@users.noreply.github.com>
Co-authored-by: Andronik Ordian <write@reusable.software>

polkadot/Cargo.lock

@@ -5044,6 +5044,7 @@ name = "polkadot-node-core-av-store"
 version = "0.1.0"
 dependencies = [
  "assert_matches",
+ "bitvec",
  "env_logger 0.8.2",
  "futures 0.3.8",
  "futures-timer 3.0.2",
@@ -5052,6 +5053,7 @@ dependencies = [
  "kvdb-rocksdb",
  "log",
  "parity-scale-codec",
+ "parking_lot 0.11.1",
  "polkadot-erasure-coding",
  "polkadot-node-subsystem",
  "polkadot-node-subsystem-test-helpers",
@@ -5060,6 +5062,7 @@ dependencies = [
  "polkadot-primitives",
  "sc-service",
  "sp-core",
+ "sp-keyring",
  "thiserror",
  "tracing",
  "tracing-futures",

polkadot/node/core/av-store/Cargo.toml

@@ -12,6 +12,7 @@ kvdb-rocksdb = "0.10.0"
 thiserror = "1.0.23"
 tracing = "0.1.22"
 tracing-futures = "0.2.4"
+bitvec = "0.17.4"
 
 parity-scale-codec = { version = "1.3.5", features = ["derive"] }
 erasure = { package = "polkadot-erasure-coding", path = "../../../erasure-coding" }
@@ -31,3 +32,5 @@ kvdb-memorydb = "0.8.0"
 sp-core = { git = "https://github.com/paritytech/substrate", branch = "master" }
 polkadot-node-subsystem-util = { path = "../../subsystem-util" }
 polkadot-node-subsystem-test-helpers = { path = "../../subsystem-test-helpers" }
+sp-keyring = { git = "https://github.com/paritytech/substrate", branch = "master" }
+parking_lot = "0.11.1"

polkadot/node/core/av-store/src/tests.rs

@@ -26,13 +26,15 @@ use futures::{
 
 use polkadot_primitives::v1::{
 	AvailableData, BlockData, CandidateDescriptor, CandidateReceipt, HeadData,
-	PersistedValidationData, PoV, Id as ParaId, CandidateHash,
+	PersistedValidationData, PoV, Id as ParaId, CandidateHash, Header, ValidatorId,
 };
 use polkadot_node_subsystem_util::TimeoutExt;
 use polkadot_subsystem::{
-	ActiveLeavesUpdate, errors::RuntimeApiError, JaegerSpan,
+	ActiveLeavesUpdate, errors::RuntimeApiError, JaegerSpan, messages::AllMessages,
 };
 use polkadot_node_subsystem_test_helpers as test_helpers;
+use sp_keyring::Sr25519Keyring;
+use parking_lot::Mutex;
 
 struct TestHarness {
 	virtual_overseer: test_helpers::TestSubsystemContextHandle<AvailabilityStoreMessage>,
@@ -60,9 +62,41 @@ impl TestCandidateBuilder {
 	}
 }
 
+#[derive(Clone)]
+struct TestClock {
+	inner: Arc<Mutex<Duration>>,
+}
+
+impl TestClock {
+	fn now(&self) -> Duration {
+		self.inner.lock().clone()
+	}
+
+	fn inc(&self, by: Duration) {
+		*self.inner.lock() += by;
+	}
+}
+
+impl Clock for TestClock {
+	fn now(&self) -> Result<Duration, Error> {
+		Ok(TestClock::now(self))
+	}
+}
+
+
+#[derive(Clone)]
 struct TestState {
 	persisted_validation_data: PersistedValidationData,
 	pruning_config: PruningConfig,
+	clock: TestClock,
+}
+
+impl TestState {
+	// pruning is only polled periodically, so we sometimes need to delay until
+	// we're sure the subsystem has done pruning.
+	async fn wait_for_pruning(&self) {
+		Delay::new(self.pruning_config.pruning_interval * 2).await
+	}
 }
 
 impl Default for TestState {
@@ -77,20 +111,26 @@ impl Default for TestState {
 		};
 
 		let pruning_config = PruningConfig {
-			keep_stored_block_for: Duration::from_secs(1),
-			keep_finalized_block_for: Duration::from_secs(2),
-			keep_finalized_chunk_for: Duration::from_secs(2),
+			keep_unavailable_for: Duration::from_secs(1),
+			keep_finalized_for: Duration::from_secs(2),
+			pruning_interval: Duration::from_millis(250),
+		};
+
+		let clock = TestClock {
+			inner: Arc::new(Mutex::new(Duration::from_secs(0))),
 		};
 
 		Self {
 			persisted_validation_data,
 			pruning_config,
+			clock,
 		}
 	}
 }
 
+
 fn test_harness<T: Future<Output=()>>(
-	pruning_config: PruningConfig,
+	state: TestState,
 	store: Arc<dyn KeyValueDB>,
 	test: impl FnOnce(TestHarness) -> T,
 ) {
@@ -109,7 +149,12 @@ fn test_harness<T: Future<Output=()>>(
 	let pool = sp_core::testing::TaskExecutor::new();
 	let (context, virtual_overseer) = test_helpers::make_subsystem_context(pool.clone());
 
-	let subsystem = AvailabilityStoreSubsystem::new_in_memory(store, pruning_config);
+	let subsystem = AvailabilityStoreSubsystem::new_in_memory(
+		store,
+		state.pruning_config.clone(),
+		Box::new(state.clock),
+	);
+
 	let subsystem = run(subsystem, context);
 
 	let test_fut = test(TestHarness {
@@ -170,11 +215,17 @@ async fn overseer_signal(
 		.expect(&format!("{:?} is more than enough for sending signals.", TIMEOUT));
 }
 
+fn with_tx(db: &Arc<impl KeyValueDB>, f: impl FnOnce(&mut DBTransaction)) {
+	let mut tx = DBTransaction::new();
+	f(&mut tx);
+	db.write(tx).unwrap();
+}
+
 #[test]
 fn runtime_api_error_does_not_stop_the_subsystem() {
 	let store = Arc::new(kvdb_memorydb::create(columns::NUM_COLUMNS));
 
-	test_harness(PruningConfig::default(), store, |test_harness| async move {
+	test_harness(TestState::default(), store, |test_harness| async move {
 		let TestHarness { mut virtual_overseer } = test_harness;
 		let new_leaf = Hash::repeat_byte(0x01);
 
@@ -218,7 +269,59 @@ fn runtime_api_error_does_not_stop_the_subsystem() {
 #[test]
 fn store_chunk_works() {
 	let store = Arc::new(kvdb_memorydb::create(columns::NUM_COLUMNS));
-	test_harness(PruningConfig::default(), store.clone(), |test_harness| async move {
+	test_harness(TestState::default(), store.clone(), |test_harness| async move {
+		let TestHarness { mut virtual_overseer } = test_harness;
+		let relay_parent = Hash::repeat_byte(32);
+		let candidate_hash = CandidateHash(Hash::repeat_byte(33));
+		let validator_index = 5;
+		let n_validators = 10;
+
+		let chunk = ErasureChunk {
+			chunk: vec![1, 2, 3],
+			index: validator_index,
+			proof: vec![vec![3, 4, 5]],
+		};
+
+		// Ensure an entry already exists. In reality this would come from watching
+		// chain events.
+		with_tx(&store, |tx| {
+			super::write_meta(tx, &candidate_hash, &CandidateMeta {
+				data_available: false,
+				chunks_stored: bitvec::bitvec![BitOrderLsb0, u8; 0; n_validators],
+				state: State::Unavailable(BETimestamp(0)),
+			});
+		});
+
+		let (tx, rx) = oneshot::channel();
+
+		let chunk_msg = AvailabilityStoreMessage::StoreChunk {
+			candidate_hash,
+			relay_parent,
+			chunk: chunk.clone(),
+			tx,
+		};
+
+		overseer_send(&mut virtual_overseer, chunk_msg.into()).await;
+		assert_eq!(rx.await.unwrap(), Ok(()));
+
+		let (tx, rx) = oneshot::channel();
+		let query_chunk = AvailabilityStoreMessage::QueryChunk(
+			candidate_hash,
+			validator_index,
+			tx,
+		);
+
+		overseer_send(&mut virtual_overseer, query_chunk.into()).await;
+
+		assert_eq!(rx.await.unwrap().unwrap(), chunk);
+	});
+}
+
+
+#[test]
+fn store_chunk_does_nothing_if_no_entry_already() {
+	let store = Arc::new(kvdb_memorydb::create(columns::NUM_COLUMNS));
+	test_harness(TestState::default(), store.clone(), |test_harness| async move {
 		let TestHarness { mut virtual_overseer } = test_harness;
 		let relay_parent = Hash::repeat_byte(32);
 		let candidate_hash = CandidateHash(Hash::repeat_byte(33));
@@ -240,19 +343,7 @@ fn store_chunk_works() {
 		};
 
 		overseer_send(&mut virtual_overseer, chunk_msg.into()).await;
-
-		assert_matches!(
-			overseer_recv(&mut virtual_overseer).await,
-			AllMessages::ChainApi(ChainApiMessage::BlockNumber(
-				hash,
-				tx,
-			)) => {
-				assert_eq!(hash, relay_parent);
-				tx.send(Ok(Some(4))).unwrap();
-			}
-		);
-
-		assert_eq!(rx.await.unwrap(), Ok(()));
+		assert_eq!(rx.await.unwrap(), Err(()));
 
 		let (tx, rx) = oneshot::channel();
 		let query_chunk = AvailabilityStoreMessage::QueryChunk(
@@ -263,7 +354,52 @@ fn store_chunk_works() {
 
 		overseer_send(&mut virtual_overseer, query_chunk.into()).await;
 
-		assert_eq!(rx.await.unwrap().unwrap(), chunk);
+		assert!(rx.await.unwrap().is_none());
+	});
+}
+
+#[test]
+fn query_chunk_checks_meta() {
+	let store = Arc::new(kvdb_memorydb::create(columns::NUM_COLUMNS));
+	test_harness(TestState::default(), store.clone(), |test_harness| async move {
+		let TestHarness { mut virtual_overseer } = test_harness;
+		let candidate_hash = CandidateHash(Hash::repeat_byte(33));
+		let validator_index = 5;
+		let n_validators = 10;
+
+		// Ensure an entry already exists. In reality this would come from watching
+		// chain events.
+		with_tx(&store, |tx| {
+			super::write_meta(tx, &candidate_hash, &CandidateMeta {
+				data_available: false,
+				chunks_stored: {
+					let mut v = bitvec::bitvec![BitOrderLsb0, u8; 0; n_validators];
+					v.set(validator_index as usize, true);
+					v
+				},
+				state: State::Unavailable(BETimestamp(0)),
+			});
+		});
+
+		let (tx, rx) = oneshot::channel();
+		let query_chunk = AvailabilityStoreMessage::QueryChunkAvailability(
+			candidate_hash,
+			validator_index,
+			tx,
+		);
+
+		overseer_send(&mut virtual_overseer, query_chunk.into()).await;
+		assert!(rx.await.unwrap());
+
+		let (tx, rx) = oneshot::channel();
+		let query_chunk = AvailabilityStoreMessage::QueryChunkAvailability(
+			candidate_hash,
+			validator_index + 1,
+			tx,
+		);
+
+		overseer_send(&mut virtual_overseer, query_chunk.into()).await;
+		assert!(!rx.await.unwrap());
 	});
 }
 
@@ -271,7 +407,7 @@ fn store_chunk_works() {
 fn store_block_works() {
 	let store = Arc::new(kvdb_memorydb::create(columns::NUM_COLUMNS));
 	let test_state = TestState::default();
-	test_harness(test_state.pruning_config.clone(), store.clone(), |test_harness| async move {
+	test_harness(test_state.clone(), store.clone(), |test_harness| async move {
 		let TestHarness { mut virtual_overseer } = test_harness;
 		let candidate_hash = CandidateHash(Hash::repeat_byte(1));
 		let validator_index = 5;
@@ -283,7 +419,7 @@ fn store_block_works() {
 
 		let available_data = AvailableData {
 			pov: Arc::new(pov),
-			validation_data: test_state.persisted_validation_data,
+			validation_data: test_state.persisted_validation_data.clone(),
 		};
 
 
@@ -319,13 +455,12 @@ fn store_block_works() {
 	});
 }
 
-
 #[test]
 fn store_pov_and_query_chunk_works() {
 	let store = Arc::new(kvdb_memorydb::create(columns::NUM_COLUMNS));
 	let test_state = TestState::default();
 
-	test_harness(test_state.pruning_config.clone(), store.clone(), |test_harness| async move {
+	test_harness(test_state.clone(), store.clone(), |test_harness| async move {
 		let TestHarness { mut virtual_overseer } = test_harness;
 		let candidate_hash = CandidateHash(Hash::repeat_byte(1));
 		let n_validators = 10;
@@ -336,11 +471,10 @@ fn store_pov_and_query_chunk_works() {
 
 		let available_data = AvailableData {
 			pov: Arc::new(pov),
-			validation_data: test_state.persisted_validation_data,
+			validation_data: test_state.persisted_validation_data.clone(),
 		};
 
-		let no_metrics = Metrics(None);
-		let chunks_expected = get_chunks(&available_data, n_validators as usize, &no_metrics).unwrap();
+		let chunks_expected = erasure::obtain_chunks_v1(n_validators as _, &available_data).unwrap();
 
 		let (tx, rx) = oneshot::channel();
 		let block_msg = AvailabilityStoreMessage::StoreAvailableData(
@@ -358,71 +492,17 @@ fn store_pov_and_query_chunk_works() {
 		for validator_index in 0..n_validators {
 			let chunk = query_chunk(&mut virtual_overseer, candidate_hash, validator_index).await.unwrap();
 
-			assert_eq!(chunk, chunks_expected[validator_index as usize]);
+			assert_eq!(chunk.chunk, chunks_expected[validator_index as usize]);
 		}
 	});
 }
 
-#[test]
-fn stored_but_not_included_chunk_is_pruned() {
-	let store = Arc::new(kvdb_memorydb::create(columns::NUM_COLUMNS));
-	let test_state = TestState::default();
-
-	test_harness(test_state.pruning_config.clone(), store.clone(), |test_harness| async move {
-		let TestHarness { mut virtual_overseer } = test_harness;
-		let candidate_hash = CandidateHash(Hash::repeat_byte(1));
-		let relay_parent = Hash::repeat_byte(2);
-		let validator_index = 5;
-
-		let chunk = ErasureChunk {
-			chunk: vec![1, 2, 3],
-			index: validator_index,
-			proof: vec![vec![3, 4, 5]],
-		};
-
-		let (tx, rx) = oneshot::channel();
-		let chunk_msg = AvailabilityStoreMessage::StoreChunk {
-			candidate_hash,
-			relay_parent,
-			chunk: chunk.clone(),
-			tx,
-		};
-
-		overseer_send(&mut virtual_overseer, chunk_msg.into()).await;
-
-		assert_matches!(
-			overseer_recv(&mut virtual_overseer).await,
-			AllMessages::ChainApi(ChainApiMessage::BlockNumber(
-				hash,
-				tx,
-			)) => {
-				assert_eq!(hash, relay_parent);
-				tx.send(Ok(Some(4))).unwrap();
-			}
-		);
-
-		rx.await.unwrap().unwrap();
-
-		// At this point data should be in the store.
-		assert_eq!(
-			query_chunk(&mut virtual_overseer, candidate_hash, validator_index).await.unwrap(),
-			chunk,
-		);
-
-		// Wait for twice as long as the stored block kept for.
-		Delay::new(test_state.pruning_config.keep_stored_block_for * 2).await;
-
-		// The block was not included by this point so it should be pruned now.
-		assert!(query_chunk(&mut virtual_overseer, candidate_hash, validator_index).await.is_none());
-	});
-}
-
 #[test]
 fn stored_but_not_included_data_is_pruned() {
 	let store = Arc::new(kvdb_memorydb::create(columns::NUM_COLUMNS));
 	let test_state = TestState::default();
 
-	test_harness(test_state.pruning_config.clone(), store.clone(), |test_harness| async move {
+	test_harness(test_state.clone(), store.clone(), |test_harness| async move {
 		let TestHarness { mut virtual_overseer } = test_harness;
 		let candidate_hash = CandidateHash(Hash::repeat_byte(1));
 		let n_validators = 10;
@@ -433,7 +513,7 @@ fn stored_but_not_included_data_is_pruned() {
 
 		let available_data = AvailableData {
 			pov: Arc::new(pov),
-			validation_data: test_state.persisted_validation_data,
+			validation_data: test_state.persisted_validation_data.clone(),
 		};
 
 		let (tx, rx) = oneshot::channel();
@@ -455,8 +535,9 @@ fn stored_but_not_included_data_is_pruned() {
 			available_data,
 		);
 
-		// Wait for twice as long as the stored block kept for.
-		Delay::new(test_state.pruning_config.keep_stored_block_for * 2).await;
+		// Wait until pruning.
+		test_state.clock.inc(test_state.pruning_config.keep_unavailable_for);
+		test_state.wait_for_pruning().await;
 
 		// The block was not included by this point so it should be pruned now.
 		assert!(query_available_data(&mut virtual_overseer, candidate_hash).await.is_none());
@@ -468,7 +549,7 @@ fn stored_data_kept_until_finalized() {
 	let store = Arc::new(kvdb_memorydb::create(columns::NUM_COLUMNS));
 	let test_state = TestState::default();
 
-	test_harness(test_state.pruning_config.clone(), store.clone(), |test_harness| async move {
+	test_harness(test_state.clone(), store.clone(), |test_harness| async move {
 		let TestHarness { mut virtual_overseer } = test_harness;
 		let n_validators = 10;
 
@@ -487,9 +568,12 @@ fn stored_data_kept_until_finalized() {
 
 		let available_data = AvailableData {
 			pov: Arc::new(pov),
-			validation_data: test_state.persisted_validation_data,
+			validation_data: test_state.persisted_validation_data.clone(),
 		};
 
+		let parent = Hash::repeat_byte(2);
+		let block_number = 10;
+
 		let (tx, rx) = oneshot::channel();
 		let block_msg = AvailabilityStoreMessage::StoreAvailableData(
 			candidate_hash,
@@ -509,29 +593,17 @@ fn stored_data_kept_until_finalized() {
 			available_data,
 		);
 
-		let new_leaf = Hash::repeat_byte(2);
-		overseer_signal(
+		let new_leaf = import_leaf(
 			&mut virtual_overseer,
-			OverseerSignal::ActiveLeaves(ActiveLeavesUpdate {
-				activated: vec![(new_leaf, Arc::new(JaegerSpan::Disabled))].into(),
-				deactivated: vec![].into(),
-			}),
+			parent,
+			block_number,
+			vec![CandidateEvent::CandidateIncluded(candidate, HeadData::default())],
+			(0..n_validators).map(|_| Sr25519Keyring::Alice.public().into()).collect(),
 		).await;
 
-		assert_matches!(
-			overseer_recv(&mut virtual_overseer).await,
-			AllMessages::RuntimeApi(RuntimeApiMessage::Request(
-				relay_parent,
-				RuntimeApiRequest::CandidateEvents(tx),
-			)) => {
-				assert_eq!(relay_parent, new_leaf);
-				tx.send(Ok(vec![
-					CandidateEvent::CandidateIncluded(candidate, HeadData::default()),
-				])).unwrap();
-			}
-		);
-
-		Delay::new(test_state.pruning_config.keep_stored_block_for * 10).await;
+		// Wait until unavailable data would definitely be pruned.
+		test_state.clock.inc(test_state.pruning_config.keep_unavailable_for * 10);
+		test_state.wait_for_pruning().await;
 
 		// At this point data should _still_ be in the store.
 		assert_eq!(
@@ -539,13 +611,18 @@ fn stored_data_kept_until_finalized() {
 			available_data,
 		);
 
+		assert!(
+			query_all_chunks(&mut virtual_overseer, candidate_hash, n_validators, true).await
+		);
+
 		overseer_signal(
 			&mut virtual_overseer,
-			OverseerSignal::BlockFinalized(new_leaf, 10)
+			OverseerSignal::BlockFinalized(new_leaf, block_number)
 		).await;
 
-		// Wait for a half of the time finalized data should be available for
-		Delay::new(test_state.pruning_config.keep_finalized_block_for / 2).await;
+		// Wait until unavailable data would definitely be pruned.
+		test_state.clock.inc(test_state.pruning_config.keep_finalized_for / 2);
+		test_state.wait_for_pruning().await;
 
 		// At this point data should _still_ be in the store.
 		assert_eq!(
@@ -553,115 +630,21 @@ fn stored_data_kept_until_finalized() {
 			available_data,
 		);
 
-		// Wait until it is should be gone.
-		Delay::new(test_state.pruning_config.keep_finalized_block_for).await;
+		assert!(
+			query_all_chunks(&mut virtual_overseer, candidate_hash, n_validators, true).await
+		);
+
+		// Wait until it definitely should be gone.
+		test_state.clock.inc(test_state.pruning_config.keep_finalized_for);
+		test_state.wait_for_pruning().await;
 
 		// At this point data should be gone from the store.
 		assert!(
 			query_available_data(&mut virtual_overseer, candidate_hash).await.is_none(),
 		);
-	});
-}
 
-#[test]
-fn stored_chunk_kept_until_finalized() {
-	let store = Arc::new(kvdb_memorydb::create(columns::NUM_COLUMNS));
-	let test_state = TestState::default();
-
-	test_harness(test_state.pruning_config.clone(), store.clone(), |test_harness| async move {
-		let TestHarness { mut virtual_overseer } = test_harness;
-		let relay_parent = Hash::repeat_byte(2);
-		let validator_index = 5;
-		let candidate = TestCandidateBuilder {
-			..Default::default()
-		}.build();
-		let candidate_hash = candidate.hash();
-
-		let chunk = ErasureChunk {
-			chunk: vec![1, 2, 3],
-			index: validator_index,
-			proof: vec![vec![3, 4, 5]],
-		};
-
-		let (tx, rx) = oneshot::channel();
-		let chunk_msg = AvailabilityStoreMessage::StoreChunk {
-			candidate_hash,
-			relay_parent,
-			chunk: chunk.clone(),
-			tx,
-		};
-
-		overseer_send(&mut virtual_overseer, chunk_msg.into()).await;
-
-		assert_matches!(
-			overseer_recv(&mut virtual_overseer).await,
-			AllMessages::ChainApi(ChainApiMessage::BlockNumber(
-				hash,
-				tx,
-			)) => {
-				assert_eq!(hash, relay_parent);
-				tx.send(Ok(Some(4))).unwrap();
-			}
-		);
-
-		rx.await.unwrap().unwrap();
-
-		// At this point data should be in the store.
-		assert_eq!(
-			query_chunk(&mut virtual_overseer, candidate_hash, validator_index).await.unwrap(),
-			chunk,
-		);
-
-		let new_leaf = Hash::repeat_byte(2);
-		overseer_signal(
-			&mut virtual_overseer,
-			OverseerSignal::ActiveLeaves(ActiveLeavesUpdate {
-				activated: vec![(new_leaf, Arc::new(JaegerSpan::Disabled))].into(),
-				deactivated: vec![].into(),
-			}),
-		).await;
-
-		assert_matches!(
-			overseer_recv(&mut virtual_overseer).await,
-			AllMessages::RuntimeApi(RuntimeApiMessage::Request(
-				relay_parent,
-				RuntimeApiRequest::CandidateEvents(tx),
-			)) => {
-				assert_eq!(relay_parent, new_leaf);
-				tx.send(Ok(vec![
-					CandidateEvent::CandidateIncluded(candidate, HeadData::default()),
-				])).unwrap();
-			}
-		);
-
-		Delay::new(test_state.pruning_config.keep_stored_block_for * 10).await;
-
-		// At this point data should _still_ be in the store.
-		assert_eq!(
-			query_chunk(&mut virtual_overseer, candidate_hash, validator_index).await.unwrap(),
-			chunk,
-		);
-
-		overseer_signal(
-			&mut virtual_overseer,
-			OverseerSignal::BlockFinalized(new_leaf, 10)
-		).await;
-
-		// Wait for a half of the time finalized data should be available for
-		Delay::new(test_state.pruning_config.keep_finalized_block_for / 2).await;
-
-		// At this point data should _still_ be in the store.
-		assert_eq!(
-			query_chunk(&mut virtual_overseer, candidate_hash, validator_index).await.unwrap(),
-			chunk,
-		);
-
-		// Wait until it is should be gone.
-		Delay::new(test_state.pruning_config.keep_finalized_chunk_for).await;
-
-		// At this point data should be gone from the store.
 		assert!(
-			query_available_data(&mut virtual_overseer, candidate_hash).await.is_none(),
+			query_all_chunks(&mut virtual_overseer, candidate_hash, n_validators, false).await
 		);
 	});
 }
@@ -671,9 +654,14 @@ fn forkfullness_works() {
 	let store = Arc::new(kvdb_memorydb::create(columns::NUM_COLUMNS));
 	let test_state = TestState::default();
 
-	test_harness(test_state.pruning_config.clone(), store.clone(), |test_harness| async move {
+	test_harness(test_state.clone(), store.clone(), |test_harness| async move {
 		let TestHarness { mut virtual_overseer } = test_harness;
 		let n_validators = 10;
+		let block_number_1 = 5;
+		let block_number_2 = 5;
+		let validators: Vec<_> = (0..n_validators).map(|_| Sr25519Keyring::Alice.public().into()).collect();
+		let parent_1 = Hash::repeat_byte(3);
+		let parent_2 = Hash::repeat_byte(4);
 
 		let pov_1 = PoV {
 			block_data: BlockData(vec![1, 2, 3]),
@@ -708,7 +696,7 @@ fn forkfullness_works() {
 
 		let available_data_2 = AvailableData {
 			pov: Arc::new(pov_2),
-			validation_data: test_state.persisted_validation_data,
+			validation_data: test_state.persisted_validation_data.clone(),
 		};
 
 		let (tx, rx) = oneshot::channel();
@@ -747,47 +735,25 @@ fn forkfullness_works() {
 			available_data_2,
 		);
 
-
-		let new_leaf_1 = Hash::repeat_byte(2);
-		let new_leaf_2 = Hash::repeat_byte(3);
-
-		overseer_signal(
+		let new_leaf_1 = import_leaf(
 			&mut virtual_overseer,
-			OverseerSignal::ActiveLeaves(ActiveLeavesUpdate {
-				activated: vec![(new_leaf_1, Arc::new(JaegerSpan::Disabled)), (new_leaf_2, Arc::new(JaegerSpan::Disabled))].into(),
-				deactivated: vec![].into(),
-			}),
+			parent_1,
+			block_number_1,
+			vec![CandidateEvent::CandidateIncluded(candidate_1, HeadData::default())],
+			validators.clone(),
 		).await;
 
-		assert_matches!(
-			overseer_recv(&mut virtual_overseer).await,
-			AllMessages::RuntimeApi(RuntimeApiMessage::Request(
-				leaf,
-				RuntimeApiRequest::CandidateEvents(tx),
-			)) => {
-				assert_eq!(leaf, new_leaf_1);
-				tx.send(Ok(vec![
-					CandidateEvent::CandidateIncluded(candidate_1, HeadData::default()),
-				])).unwrap();
-			}
-		);
-
-		assert_matches!(
-			overseer_recv(&mut virtual_overseer).await,
-			AllMessages::RuntimeApi(RuntimeApiMessage::Request(
-				leaf,
-				RuntimeApiRequest::CandidateEvents(tx),
-			)) => {
-				assert_eq!(leaf, new_leaf_2);
-				tx.send(Ok(vec![
-					CandidateEvent::CandidateIncluded(candidate_2, HeadData::default()),
-				])).unwrap();
-			}
-		);
+		let _new_leaf_2 = import_leaf(
+			&mut virtual_overseer,
+			parent_2,
+			block_number_2,
+			vec![CandidateEvent::CandidateIncluded(candidate_2, HeadData::default())],
+			validators.clone(),
+		).await;
 
 		overseer_signal(
 			&mut virtual_overseer,
-			OverseerSignal::BlockFinalized(new_leaf_1, 5)
+			OverseerSignal::BlockFinalized(new_leaf_1, block_number_1)
 		).await;
 
 		// Data of both candidates should be still present in the DB.
@@ -800,10 +766,41 @@ fn forkfullness_works() {
 			query_available_data(&mut virtual_overseer, candidate_2_hash).await.unwrap(),
 			available_data_2,
 		);
-		// Wait for longer than finalized blocks should be kept for
-		Delay::new(test_state.pruning_config.keep_finalized_block_for + Duration::from_secs(1)).await;
 
-		// Data of both candidates should be gone now.
+		assert!(
+			query_all_chunks(&mut virtual_overseer, candidate_1_hash, n_validators, true).await,
+		);
+
+		assert!(
+			query_all_chunks(&mut virtual_overseer, candidate_2_hash, n_validators, true).await,
+		);
+
+		// Candidate 2 should now be considered unavailable and will be pruned.
+		test_state.clock.inc(test_state.pruning_config.keep_unavailable_for);
+		test_state.wait_for_pruning().await;
+
+		assert_eq!(
+			query_available_data(&mut virtual_overseer, candidate_1_hash).await.unwrap(),
+			available_data_1,
+		);
+
+		assert!(
+			query_available_data(&mut virtual_overseer, candidate_2_hash).await.is_none(),
+		);
+
+		assert!(
+			query_all_chunks(&mut virtual_overseer, candidate_1_hash, n_validators, true).await,
+		);
+
+		assert!(
+			query_all_chunks(&mut virtual_overseer, candidate_2_hash, n_validators, false).await,
+		);
+
+		// Wait for longer than finalized blocks should be kept for
+		test_state.clock.inc(test_state.pruning_config.keep_finalized_for);
+		test_state.wait_for_pruning().await;
+
+		// Everything should be pruned now.
 		assert!(
 			query_available_data(&mut virtual_overseer, candidate_1_hash).await.is_none(),
 		);
@@ -811,6 +808,14 @@ fn forkfullness_works() {
 		assert!(
 			query_available_data(&mut virtual_overseer, candidate_2_hash).await.is_none(),
 		);
+
+		assert!(
+			query_all_chunks(&mut virtual_overseer, candidate_1_hash, n_validators, false).await,
+		);
+
+		assert!(
+			query_all_chunks(&mut virtual_overseer, candidate_2_hash, n_validators, false).await,
+		);
 	});
 }
 
@@ -838,3 +843,88 @@ async fn query_chunk(
 
 	rx.await.unwrap()
 }
+
+async fn query_all_chunks(
+	virtual_overseer: &mut test_helpers::TestSubsystemContextHandle<AvailabilityStoreMessage>,
+	candidate_hash: CandidateHash,
+	n_validators: u32,
+	expect_present: bool,
+) -> bool {
+	for i in 0..n_validators {
+		if query_chunk(virtual_overseer, candidate_hash, i).await.is_some() != expect_present {
+			return false
+		}
+	}
+	true
+}
+
+async fn import_leaf(
+	virtual_overseer: &mut test_helpers::TestSubsystemContextHandle<AvailabilityStoreMessage>,
+	parent_hash: Hash,
+	block_number: BlockNumber,
+	events: Vec<CandidateEvent>,
+	validators: Vec<ValidatorId>,
+) -> Hash {
+	let header = Header {
+		parent_hash,
+		number: block_number,
+		state_root: Hash::zero(),
+		extrinsics_root: Hash::zero(),
+		digest: Default::default(),
+	};
+	let new_leaf = header.hash();
+
+	overseer_signal(
+		virtual_overseer,
+		OverseerSignal::ActiveLeaves(ActiveLeavesUpdate {
+			activated: vec![(new_leaf, Arc::new(JaegerSpan::Disabled))].into(),
+			deactivated: vec![].into(),
+		}),
+	).await;
+
+	assert_matches!(
+		overseer_recv(virtual_overseer).await,
+		AllMessages::RuntimeApi(RuntimeApiMessage::Request(
+			relay_parent,
+			RuntimeApiRequest::CandidateEvents(tx),
+		)) => {
+			assert_eq!(relay_parent, new_leaf);
+			tx.send(Ok(events)).unwrap();
+		}
+	);
+
+	assert_matches!(
+		overseer_recv(virtual_overseer).await,
+		AllMessages::ChainApi(ChainApiMessage::BlockNumber(
+			relay_parent,
+			tx,
+		)) => {
+			assert_eq!(relay_parent, new_leaf);
+			tx.send(Ok(Some(block_number))).unwrap();
+		}
+	);
+
+	assert_matches!(
+		overseer_recv(virtual_overseer).await,
+		AllMessages::ChainApi(ChainApiMessage::BlockHeader(
+			relay_parent,
+			tx,
+		)) => {
+			assert_eq!(relay_parent, new_leaf);
+			tx.send(Ok(Some(header))).unwrap();
+		}
+	);
+
+	assert_matches!(
+		overseer_recv(virtual_overseer).await,
+		AllMessages::RuntimeApi(RuntimeApiMessage::Request(
+			relay_parent,
+			RuntimeApiRequest::Validators(tx),
+		)) => {
+			assert_eq!(relay_parent, parent_hash);
+			tx.send(Ok(validators)).unwrap();
+		}
+	);
+
+	new_leaf
+}

polkadot/roadmap/implementers-guide/src/node/utility/availability-store.md

@@ -9,25 +9,19 @@ The two data types:
 
 For each of these data we have pruning rules that determine how long we need to keep that data available.
 
-PoV hypothetically only need to be kept around until the block where the data was made fully available is finalized. However, disputes can revert finality, so we need to be a bit more conservative. We should keep the PoV until a block that finalized availability of it has been finalized for 1 day.
+PoV hypothetically only need to be kept around until the block where the data was made fully available is finalized. However, disputes can revert finality, so we need to be a bit more conservative and we add a delay. We should keep the PoV until a block that finalized availability of it has been finalized for 1 day + 1 hour.
 
-> TODO: arbitrary, but extracting `acceptance_period` is kind of hard here...
-
-Availability chunks need to be kept available until the dispute period for the corresponding candidate has ended. We can accomplish this by using the same criterion as the above, plus a delay. This gives us a pruning condition of the block finalizing availability of the chunk being final for 1 day + 1 hour.
-
-> TODO: again, concrete acceptance-period would be nicer here, but complicates things
+Availability chunks need to be kept available until the dispute period for the corresponding candidate has ended. We can accomplish this by using the same criterion as the above. This gives us a pruning condition of the block finalizing availability of the chunk being final for 1 day + 1 hour.
 
 There is also the case where a validator commits to make a PoV available, but the corresponding candidate is never backed. In this case, we keep the PoV available for 1 hour.
 
-> TODO: ideally would be an upper bound on how far back contextual execution is OK.
+There may be multiple competing blocks all ending the availability phase for a particular candidate. Until finality, it will be unclear which of those is actually the canonical chain, so the pruning records for PoVs and Availability chunks should keep track of all such blocks.
 
-There may be multiple competing blocks all ending the availability phase for a particular candidate. Until (and slightly beyond) finality, it will be unclear which of those is actually the canonical chain, so the pruning records for PoVs and Availability chunks should keep track of all such blocks.
-
-## Lifetime of the PoV in the storage
+## Lifetime of the block data and chunks in storage
 
 ```dot process
 digraph {
-	label = "Block life FSM\n\n\n";
+	label = "Block data FSM\n\n\n";
 	labelloc = "t";
 	rankdir="LR";
 
@@ -39,32 +33,56 @@ digraph {
 	st -> inc [label = "Block\nincluded"]
 	st -> prn [label = "Stored block\ntimed out"]
 	inc -> fin [label = "Block\nfinalized"]
-	fin -> prn [label = "Block keep time\n(1 day) elapsed"]
-}
-```
-
-## Lifetime of the chunk in the storage
-
-```dot process
-digraph {
-	label = "Chunk life FSM\n\n\n";
-	labelloc = "t";
-	rankdir="LR";
-
-	chst [label = "Chunk\nStored"; shape = circle]
-	st [label = "Block\nStored"; shape = circle]
-	inc [label = "Included"; shape = circle]
-	fin [label = "Finalized"; shape = circle]
-	prn [label = "Pruned"; shape = circle]
-
-	chst -> inc [label = "Block\nincluded"]
-	st -> inc [label = "Block\nincluded"]
-	st -> prn [label = "Stored block\ntimed out"]
-	inc -> fin [label = "Block\nfinalized"]
+	inc -> st [label = "Competing blocks\nfinalized"]
 	fin -> prn [label = "Block keep time\n(1 day + 1 hour) elapsed"]
 }
 ```
 
+## Database Schema
+
+We use an underlying Key-Value database where we assume we have the following operations available:
+  * `write(key, value)`
+  * `read(key) -> Option<value>`
+  * `iter_with_prefix(prefix) -> Iterator<(key, value)>` - gives all keys and values in lexicographical order where the key starts with `prefix`.
+
+We use this database to encode the following schema:
+
+```
+("available", CandidateHash) -> Option<AvailableData>
+("chunk", CandidateHash, u32) -> Option<ErasureChunk>
+("meta", CandidateHash) -> Option<CandidateMeta>
+
+("unfinalized", BlockNumber, BlockHash, CandidateHash) -> Option<()>
+("prune_by_time", Timestamp, CandidateHash) -> Option<()>
+```
+
+Timestamps are the wall-clock seconds since unix epoch. Timestamps and block numbers are both encoded as big-endian so lexicographic order is ascending.
+
+The meta information that we track per-candidate is defined as the `CandidateMeta` struct
+
+```rust
+struct CandidateMeta {
+  state: State,
+  data_available: bool,
+  chunks_stored: Bitfield,
+}
+
+enum State {
+  /// Candidate data was first observed at the given time but is not available in any block.
+  Unavailable(Timestamp),
+  /// The candidate was first observed at the given time and was included in the given list of unfinalized blocks, which may be
+  /// empty. The timestamp here is not used for pruning. Either one of these blocks will be finalized or the state will regress to
+  /// `State::Unavailable`, in which case the same timestamp will be reused.
+  Unfinalized(Timestamp, Vec<(BlockNumber, BlockHash)>),
+  /// Candidate data has appeared in a finalized block and did so at the given time.
+  Finalized(Timestamp)
+}
+```
+
+We maintain the invariant that if a candidate has a meta entry, its available data exists on disk if `data_available` is true. All chunks mentioned in the meta entry are available.
+
+Additionally, there is exactly one `prune_by_time` entry which holds the candidate hash unless the state is `Unfinalized`. There may be zero, one, or many "unfinalized" keys with the given candidate, and this will correspond to the `state` of the meta entry.
+
 ## Protocol
 
 Input: [`AvailabilityStoreMessage`][ASM]
@@ -72,97 +90,81 @@ Input: [`AvailabilityStoreMessage`][ASM]
 Output:
 - [`RuntimeApiMessage`][RAM]
 
+
 ## Functionality
 
-On `ActiveLeavesUpdate`:
-
 For each head in the `activated` list:
-	- Note any new candidates backed in the block. Update pruning records for any stored `PoVBlock`s.
-	- Note any newly-included candidates backed in the block. Update pruning records for any stored availability chunks.
+  - Note any new candidates backed in the block. Update the `CandidateMeta` for each. If the `CandidateMeta` does not exist, create it as `Unavailable` with the current timestamp. Register a `"prune_by_time"` entry based on the current timestamp + 1 hour.
+  - Note any new candidate included in the block. Update the `CandidateMeta` for each, performing a transition from `Unavailable` to `Unfinalized` if necessary. That includes removing the `"prune_by_time"` entry. Add the block hash and number to the state, if unfinalized. Add an `"unfinalized"` entry for the block and candidate.
+  - The `CandidateEvent` runtime API can be used for this purpose.
+  - TODO: load all ancestors of the head back to the finalized block so we don't miss anything if import notifications are missed. If a `StoreChunk` message is received for a candidate which has no entry, then we will prematurely lose the data.
 
-On `OverseerSignal::BlockFinalized(_)` events:
+On `OverseerSignal::BlockFinalized(finalized)` events:
+  - for each key in `iter_with_prefix("unfinalized")`
+    - Stop if the key is beyond `("unfinalized, finalized)`
+    - For each block number f that we encounter, load the finalized hash for that block.
+      - The state of each `CandidateMeta` we encounter here must be `Unfinalized`, since we loaded the candidate from an `"unfinalized"` key.
+      - For each candidate that we encounter under `f` and the finalized block hash,
+        - Update the `CandidateMeta` to have `State::Finalized`.  Remove all `"unfinalized"` entries from the old `Unfinalized` state.
+        - Register a `"prune_by_time"` entry for the candidate based on the current time + 1 day + 1 hour.
+      - For each candidate that we encounter under `f` which is not under the finalized block hash,
+        - Remove all entries under `f` in the `Unfinalized` state.
+        - If the `CandidateMeta` has state `Unfinalized` with an empty list of blocks, downgrade to `Unavailable` and re-schedule pruning under the timestamp + 1 hour. We do not prune here as the candidate still may be included in a descendent of the finalized chain.
+      - Remove all `"unfinalized"` keys under `f`.  
+  - Update last_finalized = finalized.
 
-- Handle all pruning based on the newly-finalized block.
+  This is roughly `O(n * m)` where n is the number of blocks finalized since the last update, and `m` is the number of parachains.
 
-On `QueryPoV` message:
+On `QueryAvailableData` message:
 
-- Return the PoV block, if any, for that candidate hash.
+  - Query `("available", candidate_hash)`
+
+  This is `O(n)` in the size of the data, which may be large.
+
+On `QueryDataAvailability` message:
+
+  - Query whether `("meta", candidate_hash)` exists and `data_available == true`.
+
+  This is `O(n)` in the size of the metadata which is small.
 
 On `QueryChunk` message:
 
-- Determine if we have the chunk indicated by the parameters and return it and its inclusion proof via the response channel if so.
+  - Query `("chunk", candidate_hash, index)`
+
+  This is `O(n)` in the size of the data, which may be large.
+
+On `QueryChunkAvailability message:
+
+  - Query whether `("meta", candidate_hash)` exists and the bit at `index` is set.
+
+  This is `O(n)` in the size of the metadata which is small.
 
 On `StoreChunk` message:
 
-- Store the chunk along with its inclusion proof under the candidate hash and validator index.
+  - If there is a `CandidateMeta` under the candidate hash, set the bit of the erasure-chunk in the `chunks_stored` bitfield to `1`. If it was not `1` already, write the chunk under `("chunk", candidate_hash, chunk_index)`.
 
-On `StorePoV` message:
+  This is `O(n)` in the size of the chunk.
 
-- Store the block, if the validator index is provided, store the respective chunk as well.
+On `StoreAvailableData` message:
 
-On finality event:
+  - If there is no `CandidateMeta` under the candidate hash, create it with `State::Unavailable(now)`. Load the `CandidateMeta` otherwise.
+  - Store `data` under `("available", candidate_hash)` and set `data_available` to true.
+  - Store each chunk under `("chunk", candidate_hash, index)` and set every bit in `chunks_stored` to `1`.
 
-- For the finalized block and any earlier block (if any) update pruning records of `PoV`s and chunks to keep them for respective periods after finality.
+  This is `O(n)` in the size of the data as the aggregate size of the chunks is proportional to the data.
 
-### Note any backed, included and timedout candidates in the block by `hash`.
+Every 5 minutes, run a pruning routine:
 
-- Create a `(sender, receiver)` pair.
-- Dispatch a [`RuntimeApiMessage`][RAM]`::Request(hash, RuntimeApiRequest::CandidateEvents(sender)` and listen on the receiver for a response.
-- For every event in the response:`CandidateEvent::CandidateIncluded`.
-  * For every `CandidateEvent::CandidateBacked` do nothing
-  * For every `CandidateEvent::CandidateIncluded` update pruning records of any blocks that the node stored previously.
-  * For every `CandidateEvent::CandidateTimedOut` use pruning records to prune the data; delete the info from records.
+  - for each key in `iter_with_prefix("prune_by_time")`:
+    - If the key is beyond ("prune_by_time", now), return.
+    - Remove the key.
+    - Extract `candidate_hash` from the key.
+    - Load and remove the `("meta", candidate_hash)`
+    - For each erasure chunk bit set, remove `("chunk", candidate_hash, bit_index)`.
+    - If `data_available`, remove `("available", candidate_hash)
 
-## Schema
-
-### PoV pruning
-
-We keep a record about every PoV we keep, tracking its state and the time after which this PoV should be pruned.
-
-As the state of the `Candidate` changes, so does the `Prune At` time according to the rules defined earlier.
-
-| Record 1       | .. | Record N       |
-|----------------|----|----------------|
-| CandidateHash1 | .. | CandidateHashN |
-|    Prune At    | .. |    Prune At    |
-| CandidateState | .. | CandidateState |
-
-### Chunk pruning
-
-Chunk pruning is organized in a similar schema as PoV pruning.
-
-| Record 1       | .. | Record N       |
-|----------------|----|----------------|
-| CandidateHash1 | .. | CandidateHashN |
-|    Prune At    | .. |    Prune At    |
-| CandidateState | .. | CandidateState |
-
-### Included blocks caching
-
-In order to process finality events correctly we need to cache the set of parablocks included into each relay block beginning with the last finalized block and up to the most recent heads. We have to cache this data since we are only able to query this info from the state for the `k` last blocks where `k` is a relatively small number (for more info see `Assumptions`)
-
-These are used to update Chunk pruning and PoV pruning records upon finality:
-When another block finality notification is received:
- - For any record older than this block:
-   - Update pruning
-   - Remove the record
-
-| Relay Block N | .. | Chain Head 1 | Chain Head 2 |
-|---------------|----|--------------|--------------|
-| CandidateN_1 Included | .. | Candidate1_1 Included | Candidate2_1 Included |
-| CandidateN_2 Included | .. | Candidate1_2 Included | Candidete2_2 Included |
-|          ..           | .. |          ..           |           ..          |
-| CandidateN_M Included | .. | Candidate1_K Included | Candidate2_L Included |
-
-> TODO: It's likely we will have to have a way to go from block hash to `BlockNumber` to make this work.
-
-### Blocks
-
-Blocks are simply stored as `(Hash, AvailableData)` key-value pairs.
-
-### Chunks
-
-Chunks are stored as `(Hash, Vec<ErasureChunk>)` key-value pairs.
+  This is O(n * m) in the amount of candidates and average size of the data stored. This is probably the most expensive operation but does not need
+  to be run very often.
 
 ## Basic scenarios to test
 

polkadot/roadmap/implementers-guide/src/types/overseer-protocol.md

@@ -173,12 +173,12 @@ enum AvailabilityStoreMessage {
 	QueryDataAvailability(CandidateHash, ResponseChannel<bool>),
 	/// Query a specific availability chunk of the candidate's erasure-coding by validator index.
 	/// Returns the chunk and its inclusion proof against the candidate's erasure-root.
-	QueryChunk(CandidateHash, ValidatorIndex, ResponseChannel<Option<AvailabilityChunkAndProof>>),
+	QueryChunk(CandidateHash, ValidatorIndex, ResponseChannel<Option<ErasureChunk>>),
 	/// Store a specific chunk of the candidate's erasure-coding, with an
 	/// accompanying proof.
-	StoreChunk(CandidateHash, AvailabilityChunkAndProof, ResponseChannel<Result<()>>),
+	StoreChunk(CandidateHash, ErasureChunk, ResponseChannel<Result<()>>),
 	/// Store `AvailableData`. If `ValidatorIndex` is provided, also store this validator's
-	/// `AvailabilityChunkAndProof`.
+	/// `ErasureChunk`.
 	StoreAvailableData(CandidateHash, Option<ValidatorIndex>, u32, AvailableData, ResponseChannel<Result<()>>),
 }
 ```