Erasure encoding availability (#345)

* Erasure encoding availability initial commit

 * Modifications to availability store to keep chunks as well as
   reconstructed blocks and extrinsics.
 * Gossip messages containig signed erasure chunks.
 * Requesting eraure chunks with polkadot-specific messages.
 * Validation of erasure chunk messages.

* Apply suggestions from code review

Co-Authored-By: Luke Schoen <ltfschoen@users.noreply.github.com>

* Fix build after a merge

* Gossip erasure chunk messages under their own topic

* erasure_chunks should use the appropriate topic

* Updates Cargo.lock

* Fixes after merge

* Removes a couple of leftover pieces of code

* Fixes simple stuff from review

* Updates erasure and storage for more flexible logic

* Changes validation and candidate receipt production.

* Adds add_erasure_chunks method

* Fixes most of the nits

* Better validate_collation and validate_receipt functions

* Fixes the tests

* Apply suggestions from code review

Co-Authored-By: Robert Habermeier <rphmeier@gmail.com>

* Removes unwrap() calls

* Removes ErasureChunks primitive

* Removes redundant fields from ErasureChunk struct

* AvailabilityStore should store CandidateReceipt

* Changes the way chunk messages are imported and validated.

 * Availability store now stores a validator_index and n_validators for
 each relay_parent.
 * Availability store now also stores candidate receipts.
 * Removes importing chunks in the table and moves it into network
 gossip validation.
 * Validation of erasure messages id done against receipts that are
 stored in the availability store.

* Correctly compute topics for erasure messages

* Removes an unused parameter

* Refactors availability db querying into a helper

* Adds the apis described in the writeup

* Adds a runtime api to extract erasure roots form raw extrinsics.

* Adds a barebone BlockImport impl for avalability store

* Adds the implementation of the availability worker

* Fix build after the merge with master.

* Make availability store API async

* Bring back the default wasmtime feature

* Lines width

* Bump runtime version

* Formatting and dead code elimination

* some style nits (#1)

* More nits and api cleanup

* Disable wasm CI for availability-store

* Another nit

* Formatting

polkadot/availability-store/src/lib.rs

@@ -14,28 +14,56 @@
 // You should have received a copy of the GNU General Public License
 // along with Polkadot.  If not, see <http://www.gnu.org/licenses/>.
 
-//! Persistent database for parachain data: PoV block data and outgoing messages.
+//! Persistent database for parachain data: PoV block data, erasure-coding chunks and outgoing messages.
 //!
 //! This will be written into during the block validation pipeline, and queried
 //! by networking code in order to circulate required data and maintain availability
 //! of it.
 
-use codec::{Encode, Decode};
-use kvdb::{KeyValueDB, DBTransaction};
-use polkadot_primitives::Hash;
-use polkadot_primitives::parachain::{Id as ParaId, BlockData, Message};
+#![warn(missing_docs)]
+
+use futures::prelude::*;
+use futures::channel::{mpsc, oneshot};
+use keystore::KeyStorePtr;
+use polkadot_primitives::{
+	Hash, Block,
+	parachain::{
+		Id as ParaId, BlockData, CandidateReceipt, Message, AvailableMessages, ErasureChunk,
+		ParachainHost,
+	},
+};
+use sp_runtime::traits::{BlakeTwo256, Hash as HashT, ProvideRuntimeApi};
+use sp_blockchain::{Result as ClientResult};
+use client::{
+	BlockchainEvents, BlockBody,
+};
+use sp_api::ApiExt;
+
 use log::warn;
 
+use std::sync::Arc;
 use std::collections::HashSet;
 use std::path::PathBuf;
-use std::sync::Arc;
 use std::io;
 
-mod columns {
-	pub const DATA: Option<u32> = Some(0);
-	pub const META: Option<u32> = Some(1);
-	pub const NUM_COLUMNS: u32 = 2;
-}
+mod worker;
+mod store;
+
+pub use worker::AvailabilityBlockImport;
+
+use worker::{
+	Worker, WorkerHandle, Chunks, ParachainBlocks, WorkerMsg, MakeAvailable,
+};
+
+use store::{Store as InnerStore};
+
+/// Abstraction over an executor that lets you spawn tasks in the background.
+pub(crate) type TaskExecutor =
+	Arc<dyn futures01::future::Executor<
+		Box<dyn futures01::Future<Item = (), Error = ()> + Send>
+	> + Send + Sync>;
+
+const LOG_TARGET: &str = "availability";
 
 /// Configuration for the availability store.
 pub struct Config {
@@ -45,67 +73,153 @@ pub struct Config {
 	pub path: PathBuf,
 }
 
+/// Compute gossip topic for the erasure chunk messages given the relay parent,
+/// root and the chunk index.
+///
+/// Since at this point we are not able to use [`network`] directly, but both
+/// of them need to compute these topics, this lives here and not there.
+///
+/// [`network`]: ../polkadot_network/index.html
+pub fn erasure_coding_topic(relay_parent: Hash, erasure_root: Hash, index: u32) -> Hash {
+	let mut v = relay_parent.as_ref().to_vec();
+	v.extend(erasure_root.as_ref());
+	v.extend(&index.to_le_bytes()[..]);
+	v.extend(b"erasure_chunks");
+
+	BlakeTwo256::hash(&v[..])
+}
+
+/// A trait that provides a shim for the [`NetworkService`] trait.
+///
+/// Currently it is not possible to use the networking code in the availability store
+/// core directly due to a number of loop dependencies it require:
+///
+/// `availability-store` -> `network` -> `availability-store`
+///
+/// `availability-store` -> `network` -> `validation` -> `availability-store`
+///
+/// So we provide this shim trait that gets implemented for a wrapper newtype in
+/// the [`network`] module.
+///
+/// [`NetworkService`]: ../polkadot_network/trait.NetworkService.html
+/// [`network`]: ../polkadot_network/index.html
+pub trait ProvideGossipMessages {
+	/// Get a stream of gossip erasure chunk messages for a given topic.
+	///
+	/// Each item is a tuple (relay_parent, candidate_hash, erasure_chunk)
+	fn gossip_messages_for(
+		&self,
+		topic: Hash,
+	) -> Box<dyn Stream<Item = (Hash, Hash, ErasureChunk)> + Send + Unpin>;
+
+	/// Gossip an erasure chunk message.
+	fn gossip_erasure_chunk(
+		&self,
+		relay_parent: Hash,
+		candidate_hash: Hash,
+		erasure_root: Hash,
+		chunk: ErasureChunk,
+	);
+}
+
 /// Some data to keep available about a parachain block candidate.
+#[derive(Debug)]
 pub struct Data {
 	/// The relay chain parent hash this should be localized to.
 	pub relay_parent: Hash,
 	/// The parachain index for this candidate.
 	pub parachain_id: ParaId,
-	/// Unique candidate receipt hash.
-	pub candidate_hash: Hash,
 	/// Block data.
 	pub block_data: BlockData,
 	/// Outgoing message queues from execution of the block, if any.
 	///
 	/// The tuple pairs the message queue root and the queue data.
-	pub outgoing_queues: Option<Vec<(Hash, Vec<Message>)>>,
-}
-
-fn block_data_key(relay_parent: &Hash, candidate_hash: &Hash) -> Vec<u8> {
-	(relay_parent, candidate_hash, 0i8).encode()
+	pub outgoing_queues: Option<AvailableMessages>,
 }
 
 /// Handle to the availability store.
+///
+/// This provides a proxying API that
+///   * in case of write operations provides async methods that send data to
+///     the background worker and resolve when that data is processed by the worker
+///   * in case of read opeartions queries the underlying storage synchronously.
 #[derive(Clone)]
 pub struct Store {
-	inner: Arc<dyn KeyValueDB>,
+	inner: InnerStore,
+	worker: Arc<WorkerHandle>,
+	to_worker: mpsc::UnboundedSender<WorkerMsg>,
 }
 
 impl Store {
-	/// Create a new `Store` with given config on disk.
-	#[cfg(not(target_os = "unknown"))]
-	pub fn new(config: Config) -> io::Result<Self> {
-		use kvdb_rocksdb::{Database, DatabaseConfig};
-		let mut db_config = DatabaseConfig::with_columns(Some(columns::NUM_COLUMNS));
+	/// Create a new `Store` with given condig on disk.
+	///
+	/// Creating a store among other things starts a background worker thread which
+	/// handles most of the write operations to the storage.
+	pub fn new<PGM>(config: Config, gossip: PGM) -> io::Result<Self>
+		where PGM: ProvideGossipMessages + Send + Sync + Clone + 'static
+	{
+		let inner = InnerStore::new(config)?;
+		let worker = Arc::new(Worker::start(inner.clone(), gossip));
+		let to_worker = worker.to_worker().clone();
 
-		if let Some(cache_size) = config.cache_size {
-			let mut memory_budget = std::collections::HashMap::new();
-			for i in 0..columns::NUM_COLUMNS {
-				memory_budget.insert(Some(i), cache_size / columns::NUM_COLUMNS as usize);
-			}
-
-			db_config.memory_budget = memory_budget;
-		}
-
-		let path = config.path.to_str().ok_or_else(|| io::Error::new(
-			io::ErrorKind::Other,
-			format!("Bad database path: {:?}", config.path),
-		))?;
-
-		let db = Database::open(&db_config, &path)?;
-
-		Ok(Store {
-			inner: Arc::new(db),
+		Ok(Self {
+			inner,
+			worker,
+			to_worker,
 		})
 	}
 
 	/// Create a new `Store` in-memory. Useful for tests.
-	pub fn new_in_memory() -> Self {
-		Store {
-			inner: Arc::new(::kvdb_memorydb::create(columns::NUM_COLUMNS)),
+	///
+	/// Creating a store among other things starts a background worker thread
+	/// which handles most of the write operations to the storage.
+	pub fn new_in_memory<PGM>(gossip: PGM) -> Self
+		where PGM: ProvideGossipMessages + Send + Sync + Clone + 'static
+	{
+		let inner = InnerStore::new_in_memory();
+		let worker = Arc::new(Worker::start(inner.clone(), gossip));
+		let to_worker = worker.to_worker().clone();
+
+		Self {
+			inner,
+			worker,
+			to_worker,
 		}
 	}
 
+	/// Obtain a [`BlockImport`] implementation to import blocks into this store.
+	///
+	/// This block import will act upon all newly imported blocks sending information
+	/// about parachain heads included in them to this `Store`'s background worker.
+	/// The user may create multiple instances of [`BlockImport`]s with this call.
+	///
+	/// [`BlockImport`]: https://substrate.dev/rustdocs/v1.0/substrate_consensus_common/trait.BlockImport.html
+	pub fn block_import<I, P>(
+		&self,
+		wrapped_block_import: I,
+		client: Arc<P>,
+		thread_pool: TaskExecutor,
+		keystore: KeyStorePtr,
+	) -> ClientResult<(AvailabilityBlockImport<I, P>)>
+	where
+		P: ProvideRuntimeApi + BlockchainEvents<Block> + BlockBody<Block> + Send + Sync + 'static,
+		P::Api: ParachainHost<Block>,
+		P::Api: ApiExt<Block, Error=sp_blockchain::Error>,
+	{
+		let to_worker = self.to_worker.clone();
+
+		let import = AvailabilityBlockImport::new(
+			self.inner.clone(),
+			client,
+			wrapped_block_import,
+			thread_pool,
+			keystore,
+			to_worker,
+		);
+
+		Ok(import)
+	}
+
 	/// Make some data available provisionally.
 	///
 	/// Validators with the responsibility of maintaining availability
@@ -117,174 +231,164 @@ impl Store {
 	/// to be present with the exception of the case where there is no message data
 	/// due to the block's invalidity. Determination of invalidity is beyond the
 	/// scope of this function.
-	pub fn make_available(&self, data: Data) -> io::Result<()> {
-		let mut tx = DBTransaction::new();
+	///
+	/// This method will send the `Data` to the background worker, allowing caller to
+	/// asynchrounously wait for the result.
+	pub async fn make_available(&self, data: Data) -> io::Result<()> {
+		let (s, r) = oneshot::channel();
+		let msg = WorkerMsg::MakeAvailable(MakeAvailable {
+			data,
+			result: s,
+		});
 
-		// note the meta key.
-		let mut v = match self.inner.get(columns::META, data.relay_parent.as_ref()) {
-			Ok(Some(raw)) => Vec::decode(&mut &raw[..]).expect("all stored data serialized correctly; qed"),
-			Ok(None) => Vec::new(),
-			Err(e) => {
-				warn!(target: "availability", "Error reading from availability store: {:?}", e);
-				Vec::new()
-			}
-		};
-
-		v.push(data.candidate_hash);
-		tx.put_vec(columns::META, &data.relay_parent[..], v.encode());
-
-		tx.put_vec(
-			columns::DATA,
-			block_data_key(&data.relay_parent, &data.candidate_hash).as_slice(),
-			data.block_data.encode()
-		);
-
-		if let Some(outgoing_queues) = data.outgoing_queues {
-			// This is kept forever and not pruned.
-			for (root, messages) in outgoing_queues {
-				tx.put_vec(
-					columns::DATA,
-					root.as_ref(),
-					messages.encode(),
-				);
-			}
+		let _ = self.to_worker.unbounded_send(msg);
 
+		if let Ok(Ok(())) = r.await {
+			Ok(())
+		} else {
+			Err(io::Error::new(io::ErrorKind::Other, format!("adding erasure chunks failed")))
 		}
 
-		self.inner.write(tx)
 	}
 
-	/// Note that a set of candidates have been included in a finalized block with given hash and parent hash.
-	pub fn candidates_finalized(&self, parent: Hash, finalized_candidates: HashSet<Hash>) -> io::Result<()> {
-		let mut tx = DBTransaction::new();
+	/// Get a set of all chunks we are waiting for grouped by
+	/// `(relay_parent, erasure_root, candidate_hash, our_id)`.
+	pub fn awaited_chunks(&self) -> Option<HashSet<(Hash, Hash, Hash, u32)>> {
+		self.inner.awaited_chunks()
+	}
 
-		let v = match self.inner.get(columns::META, &parent[..]) {
-			Ok(Some(raw)) => Vec::decode(&mut &raw[..]).expect("all stored data serialized correctly; qed"),
-			Ok(None) => Vec::new(),
-			Err(e) => {
-				warn!(target: "availability", "Error reading from availability store: {:?}", e);
-				Vec::new()
-			}
-		};
-		tx.delete(columns::META, &parent[..]);
+	/// Qery which candidates were included in the relay chain block by block's parent.
+	pub fn get_candidates_in_relay_block(&self, relay_block: &Hash) -> Option<Vec<Hash>> {
+		self.inner.get_candidates_in_relay_block(relay_block)
+	}
 
-		for candidate_hash in v {
-			if !finalized_candidates.contains(&candidate_hash) {
-				tx.delete(columns::DATA, block_data_key(&parent, &candidate_hash).as_slice());
-			}
+	/// Make a validator's index and a number of validators at a relay parent available.
+	///
+	/// This information is needed before the `add_candidates_in_relay_block` is called
+	/// since that call forms the awaited frontier of chunks.
+	/// In the current implementation this function is called in the `get_or_instantiate` at
+	/// the start of the parachain agreement process on top of some parent hash.
+	pub fn add_validator_index_and_n_validators(
+		&self,
+		relay_parent: &Hash,
+		validator_index: u32,
+		n_validators: u32,
+	) -> io::Result<()> {
+		self.inner.add_validator_index_and_n_validators(
+			relay_parent,
+			validator_index,
+			n_validators,
+		)
+	}
+
+	/// Query a validator's index and n_validators by relay parent.
+	pub fn get_validator_index_and_n_validators(&self, relay_parent: &Hash) -> Option<(u32, u32)> {
+		self.inner.get_validator_index_and_n_validators(relay_parent)
+	}
+
+	/// Adds an erasure chunk to storage.
+	///
+	/// The chunk should be checked for validity against the root of encoding
+	/// and its proof prior to calling this.
+	///
+	/// This method will send the chunk to the background worker, allowing caller to
+	/// asynchrounously wait for the result.
+	pub async fn add_erasure_chunk(
+		&self,
+		relay_parent: Hash,
+		receipt: CandidateReceipt,
+		chunk: ErasureChunk,
+	) -> io::Result<()> {
+		self.add_erasure_chunks(relay_parent, receipt, vec![chunk]).await
+	}
+
+	/// Adds a set of erasure chunks to storage.
+	///
+	/// The chunks should be checked for validity against the root of encoding
+	/// and it's proof prior to calling this.
+	///
+	/// This method will send the chunks to the background worker, allowing caller to
+	/// asynchrounously waiting for the result.
+	pub async fn add_erasure_chunks<I>(
+		&self,
+		relay_parent: Hash,
+		receipt: CandidateReceipt,
+		chunks: I,
+	) -> io::Result<()>
+		where I: IntoIterator<Item = ErasureChunk>
+	{
+		self.add_candidate(relay_parent, receipt.clone()).await?;
+		let (s, r) = oneshot::channel();
+		let chunks = chunks.into_iter().collect();
+		let candidate_hash = receipt.hash();
+		let msg = WorkerMsg::Chunks(Chunks {
+			relay_parent,
+			candidate_hash,
+			chunks,
+			result: s,
+		});
+
+		let _ = self.to_worker.unbounded_send(msg);
+
+		if let Ok(Ok(())) = r.await {
+			Ok(())
+		} else {
+			Err(io::Error::new(io::ErrorKind::Other, format!("adding erasure chunks failed")))
 		}
+	}
 
-		self.inner.write(tx)
+	/// Queries an erasure chunk by its block's parent and hash and index.
+	pub fn get_erasure_chunk(
+		&self,
+		relay_parent: &Hash,
+		block_data_hash: Hash,
+		index: usize,
+	) -> Option<ErasureChunk> {
+		self.inner.get_erasure_chunk(relay_parent, block_data_hash, index)
+	}
+
+	/// Stores a candidate receipt.
+	pub async fn add_candidate(
+		&self,
+		relay_parent: Hash,
+		receipt: CandidateReceipt,
+	) -> io::Result<()> {
+		let (s, r) = oneshot::channel();
+
+		let msg = WorkerMsg::ParachainBlocks(ParachainBlocks {
+			relay_parent,
+			blocks: vec![(receipt, None)],
+			result: s,
+		});
+
+		let _ = self.to_worker.unbounded_send(msg);
+
+		if let Ok(Ok(())) = r.await {
+			Ok(())
+		} else {
+			Err(io::Error::new(io::ErrorKind::Other, format!("adding erasure chunks failed")))
+		}
+	}
+
+	/// Queries a candidate receipt by it's hash.
+	pub fn get_candidate(&self, candidate_hash: &Hash) -> Option<CandidateReceipt> {
+		self.inner.get_candidate(candidate_hash)
 	}
 
 	/// Query block data.
-	pub fn block_data(&self, relay_parent: Hash, candidate_hash: Hash) -> Option<BlockData> {
-		let encoded_key = block_data_key(&relay_parent, &candidate_hash);
-		match self.inner.get(columns::DATA, &encoded_key[..]) {
-			Ok(Some(raw)) => Some(
-				BlockData::decode(&mut &raw[..]).expect("all stored data serialized correctly; qed")
-			),
-			Ok(None) => None,
-			Err(e) => {
-				warn!(target: "availability", "Error reading from availability store: {:?}", e);
-				None
-			}
-		}
+	pub fn block_data(&self, relay_parent: Hash, block_data_hash: Hash) -> Option<BlockData> {
+		self.inner.block_data(relay_parent, block_data_hash)
+	}
+
+	/// Query block data by corresponding candidate receipt's hash.
+	pub fn block_data_by_candidate(&self, relay_parent: Hash, candidate_hash: Hash)
+		-> Option<BlockData>
+	{
+		self.inner.block_data_by_candidate(relay_parent, candidate_hash)
 	}
 
 	/// Query message queue data by message queue root hash.
 	pub fn queue_by_root(&self, queue_root: &Hash) -> Option<Vec<Message>> {
-		match self.inner.get(columns::DATA, queue_root.as_ref()) {
-			Ok(Some(raw)) => Some(
-				<_>::decode(&mut &raw[..]).expect("all stored data serialized correctly; qed")
-			),
-			Ok(None) => None,
-			Err(e) => {
-				warn!(target: "availability", "Error reading from availability store: {:?}", e);
-				None
-			}
-		}
-	}
-}
-
-#[cfg(test)]
-mod tests {
-	use super::*;
-
-	#[test]
-	fn finalization_removes_unneeded() {
-		let relay_parent = [1; 32].into();
-
-		let para_id_1 = 5.into();
-		let para_id_2 = 6.into();
-
-		let candidate_1 = [2; 32].into();
-		let candidate_2 = [3; 32].into();
-
-		let block_data_1 = BlockData(vec![1, 2, 3]);
-		let block_data_2 = BlockData(vec![4, 5, 6]);
-
-		let store = Store::new_in_memory();
-		store.make_available(Data {
-			relay_parent,
-			parachain_id: para_id_1,
-			candidate_hash: candidate_1,
-			block_data: block_data_1.clone(),
-			outgoing_queues: None,
-		}).unwrap();
-
-		store.make_available(Data {
-			relay_parent,
-			parachain_id: para_id_2,
-			candidate_hash: candidate_2,
-			block_data: block_data_2.clone(),
-			outgoing_queues: None,
-		}).unwrap();
-
-		assert_eq!(store.block_data(relay_parent, candidate_1).unwrap(), block_data_1);
-		assert_eq!(store.block_data(relay_parent, candidate_2).unwrap(), block_data_2);
-
-		store.candidates_finalized(relay_parent, [candidate_1].iter().cloned().collect()).unwrap();
-
-		assert_eq!(store.block_data(relay_parent, candidate_1).unwrap(), block_data_1);
-		assert!(store.block_data(relay_parent, candidate_2).is_none());
-	}
-
-	#[test]
-	fn queues_available_by_queue_root() {
-		let relay_parent = [1; 32].into();
-		let para_id = 5.into();
-		let candidate = [2; 32].into();
-		let block_data = BlockData(vec![1, 2, 3]);
-
-		let message_queue_root_1 = [0x42; 32].into();
-		let message_queue_root_2 = [0x43; 32].into();
-
-		let message_a = Message(vec![1, 2, 3, 4]);
-		let message_b = Message(vec![4, 5, 6, 7]);
-
-		let outgoing_queues = vec![
-			(message_queue_root_1, vec![message_a.clone()]),
-			(message_queue_root_2, vec![message_b.clone()]),
-		];
-
-		let store = Store::new_in_memory();
-		store.make_available(Data {
-			relay_parent,
-			parachain_id: para_id,
-			candidate_hash: candidate,
-			block_data: block_data.clone(),
-			outgoing_queues: Some(outgoing_queues),
-		}).unwrap();
-
-		assert_eq!(
-			store.queue_by_root(&message_queue_root_1),
-			Some(vec![message_a]),
-		);
-
-		assert_eq!(
-			store.queue_by_root(&message_queue_root_2),
-			Some(vec![message_b]),
-		);
+		self.inner.queue_by_root(queue_root)
 	}
 }

polkadot/availability-store/src/store.rs

@@ -0,0 +1,689 @@
+// Copyright 2018 Parity Technologies (UK) Ltd.
+// This file is part of Polkadot.
+
+// Polkadot is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+
+// Polkadot is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+
+// You should have received a copy of the GNU General Public License
+// along with Polkadot.  If not, see <http://www.gnu.org/licenses/>.
+
+use kvdb_rocksdb::{Database, DatabaseConfig};
+use kvdb::{KeyValueDB, DBTransaction};
+use codec::{Encode, Decode};
+use polkadot_erasure_coding::{self as erasure};
+use polkadot_primitives::{
+	Hash,
+	parachain::{
+		BlockData, CandidateReceipt, Message, ErasureChunk
+	},
+};
+
+use log::{trace, warn};
+use std::collections::HashSet;
+use std::sync::Arc;
+use std::iter::FromIterator;
+use std::io;
+
+use crate::{LOG_TARGET, Data, Config};
+
+mod columns {
+	pub const DATA: Option<u32> = Some(0);
+	pub const META: Option<u32> = Some(1);
+	pub const NUM_COLUMNS: u32 = 2;
+}
+
+#[derive(Clone)]
+pub struct Store {
+	inner: Arc<dyn KeyValueDB>,
+}
+
+fn block_data_key(relay_parent: &Hash, block_data_hash: &Hash) -> Vec<u8> {
+	(relay_parent, block_data_hash, 0i8).encode()
+}
+
+fn erasure_chunks_key(relay_parent: &Hash, block_data_hash: &Hash) -> Vec<u8> {
+	(relay_parent, block_data_hash, 1i8).encode()
+}
+
+fn awaited_chunks_key() -> Vec<u8> {
+	"awaited_chunks_key".encode()
+}
+
+fn available_chunks_key(relay_parent: &Hash, erasure_root: &Hash) -> Vec<u8> {
+	(relay_parent, erasure_root, 2i8).encode()
+}
+
+fn block_to_candidate_key(block_data_hash: &Hash) -> Vec<u8> {
+	(block_data_hash, 1i8).encode()
+}
+
+fn candidate_key(candidate_hash: &Hash) -> Vec<u8> {
+	(candidate_hash, 2i8).encode()
+}
+
+fn validator_index_and_n_validators_key(relay_parent: &Hash) -> Vec<u8> {
+	(relay_parent, 3i8).encode()
+}
+
+fn candidates_in_relay_chain_block_key(relay_block: &Hash) -> Vec<u8> {
+	(relay_block, 4i8).encode()
+}
+
+fn erasure_roots_in_relay_chain_block_key(relay_block: &Hash) -> Vec<u8> {
+	(relay_block, 5i8).encode()
+}
+
+impl Store {
+	/// Create a new `Store` with given condig on disk.
+	pub(super) fn new(config: Config) -> io::Result<Self> {
+		let mut db_config = DatabaseConfig::with_columns(Some(columns::NUM_COLUMNS));
+
+		if let Some(cache_size) = config.cache_size {
+			let mut memory_budget = std::collections::HashMap::new();
+			for i in 0..columns::NUM_COLUMNS {
+				memory_budget.insert(Some(i), cache_size / columns::NUM_COLUMNS as usize);
+			}
+
+			db_config.memory_budget = memory_budget;
+		}
+
+		let path = config.path.to_str().ok_or_else(|| io::Error::new(
+			io::ErrorKind::Other,
+			format!("Bad database path: {:?}", config.path),
+		))?;
+
+		let db = Database::open(&db_config, &path)?;
+
+		Ok(Store {
+			inner: Arc::new(db),
+		})
+	}
+
+	/// Create a new `Store` in-memory. Useful for tests.
+	pub(super) fn new_in_memory() -> Self {
+		Store {
+			inner: Arc::new(::kvdb_memorydb::create(columns::NUM_COLUMNS)),
+		}
+	}
+
+	/// Make some data available provisionally.
+	pub(crate) fn make_available(&self, data: Data) -> io::Result<()> {
+		let mut tx = DBTransaction::new();
+
+		// note the meta key.
+		let mut v = self.query_inner(columns::META, data.relay_parent.as_ref()).unwrap_or(Vec::new());
+		v.push(data.block_data.hash());
+		tx.put_vec(columns::META, &data.relay_parent[..], v.encode());
+
+		tx.put_vec(
+			columns::DATA,
+			block_data_key(&data.relay_parent, &data.block_data.hash()).as_slice(),
+			data.block_data.encode()
+		);
+
+		if let Some(outgoing_queues) = data.outgoing_queues {
+			// This is kept forever and not pruned.
+			for (root, messages) in outgoing_queues.0 {
+				tx.put_vec(
+					columns::DATA,
+					root.as_ref(),
+					messages.encode(),
+				);
+			}
+
+		}
+
+		self.inner.write(tx)
+	}
+
+	/// Get a set of all chunks we are waiting for grouped by
+	/// `(relay_parent, erasure_root, candidate_hash, our_id)`.
+	pub fn awaited_chunks(&self) -> Option<HashSet<(Hash, Hash, Hash, u32)>> {
+		self.query_inner(columns::META, &awaited_chunks_key()).map(|vec: Vec<(Hash, Hash, Hash, u32)>| {
+			HashSet::from_iter(vec.into_iter())
+		})
+	}
+
+	/// Adds a set of candidates hashes that were included in a relay block by the block's parent.
+	///
+	/// If we already possess the receipts for these candidates _and_ our position at the specified
+	/// relay chain the awaited frontier of the erasure chunks will also be extended.
+	///
+	/// This method modifies the erasure chunks awaited frontier by adding this validator's
+	/// chunks from `candidates` to it. In order to do so the information about this validator's
+	/// position at parent `relay_parent` should be known to the store prior to calling this
+	/// method, in other words `add_validator_index_and_n_validators` should be called for
+	/// the given `relay_parent` before calling this function.
+	pub(crate) fn add_candidates_in_relay_block(
+		&self,
+		relay_parent: &Hash,
+		candidates: Vec<Hash>,
+	) -> io::Result<()> {
+		let mut tx = DBTransaction::new();
+		let dbkey = candidates_in_relay_chain_block_key(relay_parent);
+
+		if let Some((validator_index, _)) = self.get_validator_index_and_n_validators(relay_parent) {
+			let candidates = candidates.clone();
+			let awaited_frontier: Vec<(Hash, Hash, Hash, u32)> = self
+				.query_inner(columns::META, &awaited_chunks_key())
+				.unwrap_or_else(|| Vec::new());
+
+			let mut awaited_frontier: HashSet<(Hash, Hash, Hash, u32)> =
+				HashSet::from_iter(awaited_frontier.into_iter());
+
+			awaited_frontier.extend(candidates.into_iter().filter_map(|candidate| {
+				self.get_candidate(&candidate)
+					.map(|receipt| (relay_parent.clone(), receipt.erasure_root, candidate, validator_index))
+			}));
+			let awaited_frontier = Vec::from_iter(awaited_frontier.into_iter());
+			tx.put_vec(columns::META, &awaited_chunks_key(), awaited_frontier.encode());
+		}
+		tx.put_vec(columns::DATA, &dbkey, candidates.encode());
+
+		self.inner.write(tx)
+	}
+
+	/// Qery which candidates were included in the relay chain block by block's parent.
+	pub fn get_candidates_in_relay_block(&self, relay_block: &Hash) -> Option<Vec<Hash>> {
+		let dbkey = candidates_in_relay_chain_block_key(relay_block);
+
+		self.query_inner(columns::DATA, &dbkey)
+	}
+
+	/// Adds a set of erasure chunk roots that were included in a relay block by block's parent.
+	pub(crate) fn add_erasure_roots_in_relay_block(
+		&self,
+		relay_parent: &Hash,
+		erasure_roots: Vec<Hash>,
+	) -> io::Result<()> {
+		let mut tx = DBTransaction::new();
+		let dbkey = erasure_roots_in_relay_chain_block_key(relay_parent);
+
+		tx.put_vec(columns::DATA, &dbkey, erasure_roots.encode());
+
+		self.inner.write(tx)
+	}
+
+	/// Make a validator's index and a number of validators at a relay parent available.
+	pub(crate) fn add_validator_index_and_n_validators(
+		&self,
+		relay_parent: &Hash,
+		validator_index: u32,
+		n_validators: u32,
+	) -> io::Result<()> {
+		let mut tx = DBTransaction::new();
+		let dbkey = validator_index_and_n_validators_key(relay_parent);
+
+		tx.put_vec(columns::META, &dbkey, (validator_index, n_validators).encode());
+
+		self.inner.write(tx)
+	}
+
+	/// Query a validator's index and n_validators by relay parent.
+	pub fn get_validator_index_and_n_validators(&self, relay_parent: &Hash) -> Option<(u32, u32)> {
+		let dbkey = validator_index_and_n_validators_key(relay_parent);
+
+		self.query_inner(columns::META, &dbkey)
+	}
+
+	/// Add a set of chunks.
+	///
+	/// The same as `add_erasure_chunk` but adds a set of chunks in one atomic transaction.
+	/// Checks that all chunks have the same `relay_parent`, `block_data_hash` and `parachain_id` fields.
+	pub fn add_erasure_chunks<I>(
+		&self,
+		n_validators: u32,
+		relay_parent: &Hash,
+		candidate_hash: &Hash,
+		chunks: I,
+	) -> io::Result<()>
+		where I: IntoIterator<Item = ErasureChunk>
+	{
+		if let Some(receipt) = self.get_candidate(candidate_hash) {
+			let mut tx = DBTransaction::new();
+			let dbkey = erasure_chunks_key(relay_parent, &receipt.block_data_hash);
+
+			let mut v = self.query_inner(columns::DATA, &dbkey).unwrap_or(Vec::new());
+
+			let av_chunks_key = available_chunks_key(relay_parent, &receipt.erasure_root);
+			let mut have_chunks = self.query_inner(columns::META, &av_chunks_key).unwrap_or(Vec::new());
+
+			let awaited_frontier: Option<Vec<(Hash, Hash, Hash, u32)>> = self.query_inner(
+				columns::META,
+				&awaited_chunks_key()
+			);
+
+			for chunk in chunks.into_iter() {
+				if !have_chunks.contains(&chunk.index) {
+					have_chunks.push(chunk.index);
+				}
+				v.push(chunk);
+			}
+
+			if let Some(mut awaited_frontier) = awaited_frontier {
+				awaited_frontier.retain(|&(p, r, c, index)| {
+					!(
+						*relay_parent == p &&
+						r == receipt.erasure_root &&
+						c == receipt.hash() &&
+						have_chunks.contains(&index)
+					)
+				});
+				tx.put_vec(columns::META, &awaited_chunks_key(), awaited_frontier.encode());
+			}
+
+			// If therea are no block data and messages in the store at this point,
+			// check that they can be reconstructed now and add them to store if they can.
+			if let Ok(None) = self.inner.get(
+				columns::DATA,
+				&block_data_key(&relay_parent, &receipt.block_data_hash)
+			) {
+				if let Ok((block_data, outgoing_queues)) = erasure::reconstruct(
+					n_validators as usize,
+					v.iter().map(|chunk| (chunk.chunk.as_ref(), chunk.index as usize))) {
+					self.make_available(Data {
+						relay_parent: *relay_parent,
+						parachain_id: receipt.parachain_index,
+						block_data,
+						outgoing_queues,
+					})?;
+				}
+			}
+
+			tx.put_vec(columns::DATA, &dbkey, v.encode());
+			tx.put_vec(columns::META, &av_chunks_key, have_chunks.encode());
+
+			self.inner.write(tx)
+		} else {
+			trace!(target: LOG_TARGET, "Candidate with hash {} not found", candidate_hash);
+			Ok(())
+		}
+	}
+
+	/// Queries an erasure chunk by its block's parent and hash and index.
+	pub fn get_erasure_chunk(
+		&self,
+		relay_parent: &Hash,
+		block_data_hash: Hash,
+		index: usize,
+	) -> Option<ErasureChunk> {
+		self.query_inner(columns::DATA, &erasure_chunks_key(&relay_parent, &block_data_hash))
+			.and_then(|chunks: Vec<ErasureChunk>| {
+				chunks.iter()
+				.find(|chunk: &&ErasureChunk| chunk.index == index as u32)
+				.map(|chunk| chunk.clone())
+			})
+	}
+
+	/// Stores a candidate receipt.
+	pub fn add_candidate(&self, receipt: &CandidateReceipt) -> io::Result<()> {
+		let dbkey = candidate_key(&receipt.hash());
+		let mut tx = DBTransaction::new();
+
+		tx.put_vec(columns::DATA, &dbkey, receipt.encode());
+		tx.put_vec(columns::META, &block_to_candidate_key(&receipt.block_data_hash), receipt.hash().encode());
+
+		self.inner.write(tx)
+	}
+
+	/// Queries a candidate receipt by it's hash.
+	pub fn get_candidate(&self, candidate_hash: &Hash) -> Option<CandidateReceipt> {
+		self.query_inner(columns::DATA, &candidate_key(candidate_hash))
+	}
+
+	/// Note that a set of candidates have been included in a finalized block with given hash and parent hash.
+	pub fn candidates_finalized(
+		&self,
+		parent: Hash,
+		finalized_candidates: HashSet<Hash>,
+	) -> io::Result<()> {
+		let mut tx = DBTransaction::new();
+
+		let v = self.query_inner(columns::META, &parent[..]).unwrap_or(Vec::new());
+		tx.delete(columns::META, &parent[..]);
+
+		let awaited_frontier: Option<Vec<(Hash, Hash, Hash, u32)>> = self
+			.query_inner(columns::META, &awaited_chunks_key());
+
+		if let Some(mut awaited_frontier) = awaited_frontier {
+			awaited_frontier.retain(|&(p, c, _, _)| (p != parent && !finalized_candidates.contains(&c)));
+			tx.put_vec(columns::META, &awaited_chunks_key(), awaited_frontier.encode());
+		}
+
+		for block_data_hash in v {
+			if let Some(candidate_hash) = self.block_hash_to_candidate_hash(block_data_hash) {
+				if !finalized_candidates.contains(&candidate_hash) {
+					tx.delete(columns::DATA, block_data_key(&parent, &block_data_hash).as_slice());
+					tx.delete(columns::DATA, &erasure_chunks_key(&parent, &block_data_hash));
+					tx.delete(columns::DATA, &candidate_key(&candidate_hash));
+					tx.delete(columns::META, &block_to_candidate_key(&block_data_hash));
+				}
+			}
+		}
+
+		self.inner.write(tx)
+	}
+
+	/// Query block data.
+	pub fn block_data(&self, relay_parent: Hash, block_data_hash: Hash) -> Option<BlockData> {
+		self.query_inner(columns::DATA, &block_data_key(&relay_parent, &block_data_hash))
+	}
+
+	/// Query block data by corresponding candidate receipt's hash.
+	pub fn block_data_by_candidate(&self, relay_parent: Hash, candidate_hash: Hash) -> Option<BlockData> {
+		let receipt_key = candidate_key(&candidate_hash);
+
+		self.query_inner(columns::DATA, &receipt_key[..]).and_then(|receipt: CandidateReceipt| {
+			self.block_data(relay_parent, receipt.block_data_hash)
+		})
+	}
+
+	/// Query message queue data by message queue root hash.
+	pub fn queue_by_root(&self, queue_root: &Hash) -> Option<Vec<Message>> {
+		self.query_inner(columns::DATA, queue_root.as_ref())
+	}
+
+	fn block_hash_to_candidate_hash(&self, block_hash: Hash) -> Option<Hash> {
+		self.query_inner(columns::META, &block_to_candidate_key(&block_hash))
+	}
+
+	fn query_inner<T: Decode>(&self, column: Option<u32>, key: &[u8]) -> Option<T> {
+		match self.inner.get(column, key) {
+			Ok(Some(raw)) => {
+				let res = T::decode(&mut &raw[..]).expect("all stored data serialized correctly; qed");
+				Some(res)
+			}
+			Ok(None) => None,
+			Err(e) => {
+				warn!(target: LOG_TARGET, "Error reading from the availability store: {:?}", e);
+				None
+			}
+		}
+	}
+}
+
+#[cfg(test)]
+mod tests {
+	use super::*;
+	use polkadot_erasure_coding::{self as erasure};
+	use polkadot_primitives::parachain::{Id as ParaId, AvailableMessages};
+
+	#[test]
+	fn finalization_removes_unneeded() {
+		let relay_parent = [1; 32].into();
+
+		let para_id_1 = 5.into();
+		let para_id_2 = 6.into();
+
+		let block_data_1 = BlockData(vec![1, 2, 3]);
+		let block_data_2 = BlockData(vec![4, 5, 6]);
+
+		let erasure_chunk_1 = ErasureChunk {
+			chunk: vec![10, 20, 30],
+			index: 1,
+			proof: vec![],
+		};
+
+		let erasure_chunk_2 = ErasureChunk {
+			chunk: vec![40, 50, 60],
+			index: 1,
+			proof: vec![],
+		};
+
+		let store = Store::new_in_memory();
+		store.make_available(Data {
+			relay_parent,
+			parachain_id: para_id_1,
+			block_data: block_data_1.clone(),
+			outgoing_queues: None,
+		}).unwrap();
+
+		store.make_available(Data {
+			relay_parent,
+			parachain_id: para_id_2,
+			block_data: block_data_2.clone(),
+			outgoing_queues: None,
+		}).unwrap();
+
+		let candidate_1 = CandidateReceipt {
+			parachain_index: para_id_1,
+			collator: Default::default(),
+			signature: Default::default(),
+			head_data: Default::default(),
+			egress_queue_roots: Vec::new(),
+			fees: 0,
+			block_data_hash: block_data_1.hash(),
+			upward_messages: Vec::new(),
+			erasure_root: [6; 32].into(),
+		};
+
+		let candidate_2 = CandidateReceipt {
+			parachain_index: para_id_2,
+			collator: Default::default(),
+			signature: Default::default(),
+			head_data: Default::default(),
+			egress_queue_roots: Vec::new(),
+			fees: 0,
+			block_data_hash: block_data_2.hash(),
+			upward_messages: Vec::new(),
+			erasure_root: [6; 32].into(),
+		};
+
+		store.add_candidate(&candidate_1).unwrap();
+		store.add_candidate(&candidate_2).unwrap();
+
+		assert!(store.add_erasure_chunks(3, &relay_parent, &candidate_1.hash(), vec![erasure_chunk_1.clone()]).is_ok());
+		assert!(store.add_erasure_chunks(3, &relay_parent, &candidate_2.hash(), vec![erasure_chunk_2.clone()]).is_ok());
+
+		assert_eq!(store.block_data(relay_parent, block_data_1.hash()).unwrap(), block_data_1);
+		assert_eq!(store.block_data(relay_parent, block_data_2.hash()).unwrap(), block_data_2);
+
+		assert_eq!(store.get_erasure_chunk(&relay_parent, block_data_1.hash(), 1).as_ref(), Some(&erasure_chunk_1));
+		assert_eq!(store.get_erasure_chunk(&relay_parent, block_data_2.hash(), 1), Some(erasure_chunk_2));
+
+		assert_eq!(store.get_candidate(&candidate_1.hash()), Some(candidate_1.clone()));
+		assert_eq!(store.get_candidate(&candidate_2.hash()), Some(candidate_2.clone()));
+
+		assert_eq!(store.block_data_by_candidate(relay_parent, candidate_1.hash()).unwrap(), block_data_1);
+		assert_eq!(store.block_data_by_candidate(relay_parent, candidate_2.hash()).unwrap(), block_data_2);
+
+		store.candidates_finalized(relay_parent, [candidate_1.hash()].iter().cloned().collect()).unwrap();
+
+		assert_eq!(store.get_erasure_chunk(&relay_parent, block_data_1.hash(), 1).as_ref(), Some(&erasure_chunk_1));
+		assert!(store.get_erasure_chunk(&relay_parent, block_data_2.hash(), 1).is_none());
+
+		assert_eq!(store.get_candidate(&candidate_1.hash()), Some(candidate_1));
+		assert_eq!(store.get_candidate(&candidate_2.hash()), None);
+
+		assert_eq!(store.block_data(relay_parent, block_data_1.hash()).unwrap(), block_data_1);
+		assert!(store.block_data(relay_parent, block_data_2.hash()).is_none());
+	}
+
+	#[test]
+	fn queues_available_by_queue_root() {
+		let relay_parent = [1; 32].into();
+		let para_id = 5.into();
+		let block_data = BlockData(vec![1, 2, 3]);
+
+		let message_queue_root_1 = [0x42; 32].into();
+		let message_queue_root_2 = [0x43; 32].into();
+
+		let message_a = Message(vec![1, 2, 3, 4]);
+		let message_b = Message(vec![4, 5, 6, 7]);
+
+		let outgoing_queues = AvailableMessages(vec![
+			(message_queue_root_1, vec![message_a.clone()]),
+			(message_queue_root_2, vec![message_b.clone()]),
+		]);
+
+		let store = Store::new_in_memory();
+		store.make_available(Data {
+			relay_parent,
+			parachain_id: para_id,
+			block_data: block_data.clone(),
+			outgoing_queues: Some(outgoing_queues),
+		}).unwrap();
+
+		assert_eq!(
+			store.queue_by_root(&message_queue_root_1),
+			Some(vec![message_a]),
+		);
+
+		assert_eq!(
+			store.queue_by_root(&message_queue_root_2),
+			Some(vec![message_b]),
+		);
+	}
+
+	#[test]
+	fn erasure_coding() {
+		let relay_parent: Hash = [1; 32].into();
+		let para_id: ParaId = 5.into();
+		let block_data = BlockData(vec![42; 8]);
+		let block_data_hash = block_data.hash();
+		let n_validators = 5;
+
+		let message_queue_root_1 = [0x42; 32].into();
+		let message_queue_root_2 = [0x43; 32].into();
+
+		let message_a = Message(vec![1, 2, 3, 4]);
+		let message_b = Message(vec![5, 6, 7, 8]);
+
+		let outgoing_queues = Some(AvailableMessages(vec![
+				(message_queue_root_1, vec![message_a.clone()]),
+				(message_queue_root_2, vec![message_b.clone()]),
+		]));
+
+		let erasure_chunks = erasure::obtain_chunks(
+			n_validators,
+			&block_data,
+			outgoing_queues.as_ref()).unwrap();
+
+		let branches = erasure::branches(erasure_chunks.as_ref());
+
+		let candidate = CandidateReceipt {
+			parachain_index: para_id,
+			collator: Default::default(),
+			signature: Default::default(),
+			head_data: Default::default(),
+			egress_queue_roots: Vec::new(),
+			fees: 0,
+			block_data_hash: block_data.hash(),
+			upward_messages: Vec::new(),
+			erasure_root: [6; 32].into(),
+		};
+
+		let chunks: Vec<_> = erasure_chunks
+			.iter()
+			.zip(branches.map(|(proof, _)| proof))
+			.enumerate()
+			.map(|(index, (chunk, proof))| ErasureChunk {
+				chunk: chunk.clone(),
+				proof,
+				index: index as u32,
+			})
+			.collect();
+
+		let store = Store::new_in_memory();
+
+		store.add_candidate(&candidate).unwrap();
+		store.add_erasure_chunks(n_validators as u32, &relay_parent, &candidate.hash(), vec![chunks[0].clone()]).unwrap();
+		assert_eq!(store.get_erasure_chunk(&relay_parent, block_data_hash, 0), Some(chunks[0].clone()));
+
+		assert!(store.block_data(relay_parent, block_data_hash).is_none());
+
+		store.add_erasure_chunks(n_validators as u32, &relay_parent, &candidate.hash(), chunks).unwrap();
+		assert_eq!(store.block_data(relay_parent, block_data_hash), Some(block_data));
+	}
+
+	#[test]
+	fn add_validator_index_works() {
+		let relay_parent = [42; 32].into();
+		let store = Store::new_in_memory();
+
+		store.add_validator_index_and_n_validators(&relay_parent, 42, 24).unwrap();
+		assert_eq!(store.get_validator_index_and_n_validators(&relay_parent).unwrap(), (42, 24));
+	}
+
+	#[test]
+	fn add_candidates_in_relay_block_works() {
+		let relay_parent = [42; 32].into();
+		let store = Store::new_in_memory();
+
+		let candidates = vec![[1; 32].into(), [2; 32].into(), [3; 32].into()];
+
+		store.add_candidates_in_relay_block(&relay_parent, candidates.clone()).unwrap();
+		assert_eq!(store.get_candidates_in_relay_block(&relay_parent).unwrap(), candidates);
+	}
+
+	#[test]
+	fn awaited_chunks_works() {
+		use std::iter::FromIterator;
+		let validator_index = 3;
+		let n_validators = 10;
+		let relay_parent = [42; 32].into();
+		let erasure_root_1 = [11; 32].into();
+		let erasure_root_2 = [12; 32].into();
+		let mut receipt_1 = CandidateReceipt::default();
+		let mut receipt_2 = CandidateReceipt::default();
+
+		receipt_1.parachain_index = 1.into();
+		receipt_1.erasure_root = erasure_root_1;
+		receipt_2.parachain_index = 2.into();
+		receipt_2.erasure_root = erasure_root_2;
+
+		let chunk = ErasureChunk {
+			chunk: vec![1, 2, 3],
+			index: validator_index,
+			proof: Vec::new(),
+		};
+		let candidates = vec![receipt_1.hash(), receipt_2.hash()];
+		let erasure_roots = vec![erasure_root_1, erasure_root_2];
+
+		let store = Store::new_in_memory();
+
+		store.add_validator_index_and_n_validators(
+			&relay_parent,
+			validator_index,
+			n_validators
+		).unwrap();
+		store.add_candidate(&receipt_1).unwrap();
+		store.add_candidate(&receipt_2).unwrap();
+
+		// We are waiting for chunks from two candidates.
+		store.add_candidates_in_relay_block(&relay_parent, candidates.clone()).unwrap();
+
+		let awaited_frontier = store.awaited_chunks().unwrap();
+		warn!(target: "availability", "awaited {:?}", awaited_frontier);
+		let expected: HashSet<_> = candidates
+			.clone()
+			.into_iter()
+			.zip(erasure_roots.iter())
+			.map(|(c, e)| (relay_parent, *e, c, validator_index))
+			.collect();
+		assert_eq!(awaited_frontier, expected);
+
+		// We add chunk from one of the candidates.
+		store.add_erasure_chunks(n_validators, &relay_parent, &receipt_1.hash(), vec![chunk]).unwrap();
+
+		let awaited_frontier = store.awaited_chunks().unwrap();
+		// Now we wait for the other chunk that we haven't received yet.
+		let expected: HashSet<_> = vec![
+			(relay_parent, erasure_roots[1], candidates[1], validator_index)
+		].into_iter().collect();
+
+		assert_eq!(awaited_frontier, expected);
+
+		// Finalizing removes awaited candidates from frontier.
+		store.candidates_finalized(relay_parent, HashSet::from_iter(candidates.into_iter())).unwrap();
+
+		assert_eq!(store.awaited_chunks().unwrap().len(), 0);
+	}
+}

polkadot/availability-store/src/worker.rs

@@ -0,0 +1,998 @@
+// Copyright 2018 Parity Technologies (UK) Ltd.
+// This file is part of Polkadot.
+
+// Polkadot is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+
+// Polkadot is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+
+// You should have received a copy of the GNU General Public License
+// along with Polkadot.  If not, see <http://www.gnu.org/licenses/>.
+
+use std::collections::HashMap;
+use std::io;
+use std::sync::Arc;
+use std::thread;
+
+use log::{error, info, trace, warn};
+use sp_blockchain::{Result as ClientResult};
+use sp_runtime::traits::{Header as HeaderT, ProvideRuntimeApi};
+use sp_api::ApiExt;
+use client::{
+	BlockchainEvents, BlockBody,
+	blockchain::ProvideCache,
+};
+use consensus_common::{
+	self, BlockImport, BlockCheckParams, BlockImportParams, Error as ConsensusError,
+	ImportResult,
+	import_queue::CacheKeyId,
+};
+use polkadot_primitives::{Block, BlockId, Hash};
+use polkadot_primitives::parachain::{
+	CandidateReceipt, ParachainHost, ValidatorId,
+	ValidatorPair, AvailableMessages, BlockData, ErasureChunk,
+};
+use futures01::Future;
+use futures::channel::{mpsc, oneshot};
+use futures::{FutureExt, Sink, SinkExt, TryFutureExt, StreamExt};
+use keystore::KeyStorePtr;
+
+use tokio::runtime::current_thread::{Handle, Runtime as LocalRuntime};
+
+use crate::{LOG_TARGET, Data, TaskExecutor, ProvideGossipMessages, erasure_coding_topic};
+use crate::store::Store;
+
+/// Errors that may occur.
+#[derive(Debug, derive_more::Display, derive_more::From)]
+pub(crate) enum Error {
+	#[from]
+	StoreError(io::Error),
+	#[display(fmt = "Validator's id and number of validators at block with parent {} not found", relay_parent)]
+	IdAndNValidatorsNotFound { relay_parent: Hash },
+	#[display(fmt = "Candidate receipt with hash {} not found", candidate_hash)]
+	CandidateNotFound { candidate_hash: Hash },
+}
+
+/// Messages sent to the `Worker`.
+///
+/// Messages are sent in a number of different scenarios,
+/// for instance, when:
+///   * importing blocks in `BlockImport` implementation,
+///   * recieving finality notifications,
+///   * when the `Store` api is used by outside code.
+#[derive(Debug)]
+pub(crate) enum WorkerMsg {
+	ErasureRoots(ErasureRoots),
+	ParachainBlocks(ParachainBlocks),
+	ListenForChunks(ListenForChunks),
+	Chunks(Chunks),
+	CandidatesFinalized(CandidatesFinalized),
+	MakeAvailable(MakeAvailable),
+}
+
+/// The erasure roots of the heads included in the block with a given parent.
+#[derive(Debug)]
+pub(crate) struct ErasureRoots {
+	/// The relay parent of the block these roots belong to.
+	pub relay_parent: Hash,
+	/// The roots themselves.
+	pub erasure_roots: Vec<Hash>,
+	/// A sender to signal the result asynchronously.
+	pub result: oneshot::Sender<Result<(), Error>>,
+}
+
+/// The receipts of the heads included into the block with a given parent.
+#[derive(Debug)]
+pub(crate) struct ParachainBlocks {
+	/// The relay parent of the block these parachain blocks belong to.
+	pub relay_parent: Hash,
+	/// The blocks themselves.
+	pub blocks: Vec<(CandidateReceipt, Option<(BlockData, AvailableMessages)>)>,
+	/// A sender to signal the result asynchronously.
+	pub result: oneshot::Sender<Result<(), Error>>,
+}
+
+/// Listen gossip for these chunks.
+#[derive(Debug)]
+pub(crate) struct ListenForChunks {
+	/// The relay parent of the block the chunks from we want to listen to.
+	pub relay_parent: Hash,
+	/// The hash of the candidate chunk belongs to.
+	pub candidate_hash: Hash,
+	/// The index of the chunk we need.
+	pub index: u32,
+	/// A sender to signal the result asynchronously.
+	pub result: Option<oneshot::Sender<Result<(), Error>>>,
+}
+
+/// We have received some chunks.
+#[derive(Debug)]
+pub(crate) struct Chunks {
+	/// The relay parent of the block these chunks belong to.
+	pub relay_parent: Hash,
+	/// The hash of the parachain candidate these chunks belong to.
+	pub candidate_hash: Hash,
+	/// The chunks.
+	pub chunks: Vec<ErasureChunk>,
+	/// A sender to signal the result asynchronously.
+	pub result: oneshot::Sender<Result<(), Error>>,
+}
+
+/// These candidates have been finalized, so unneded availability may be now pruned
+#[derive(Debug)]
+pub(crate) struct CandidatesFinalized {
+	/// The relay parent of the block that was finalized.
+	relay_parent: Hash,
+	/// The parachain heads that were finalized in this block.
+	candidate_hashes: Vec<Hash>,
+}
+
+/// The message that corresponds to `make_available` call of the crate API.
+#[derive(Debug)]
+pub(crate) struct MakeAvailable {
+	/// The data being made available.
+	pub data: Data,
+	/// A sender to signal the result asynchronously.
+	pub result: oneshot::Sender<Result<(), Error>>,
+}
+
+/// An availability worker with it's inner state.
+pub(super) struct Worker<PGM> {
+	availability_store: Store,
+	provide_gossip_messages: PGM,
+	registered_gossip_streams: HashMap<Hash, exit_future::Signal>,
+
+	sender: mpsc::UnboundedSender<WorkerMsg>,
+}
+
+/// The handle to the `Worker`.
+pub(super) struct WorkerHandle {
+	exit_signal: Option<exit_future::Signal>,
+	thread: Option<thread::JoinHandle<io::Result<()>>>,
+	sender: mpsc::UnboundedSender<WorkerMsg>,
+}
+
+impl WorkerHandle {
+	pub(crate) fn to_worker(&self) -> &mpsc::UnboundedSender<WorkerMsg> {
+		&self.sender
+	}
+}
+
+impl Drop for WorkerHandle {
+	fn drop(&mut self) {
+		if let Some(signal) = self.exit_signal.take() {
+			signal.fire();
+		}
+
+		if let Some(thread) = self.thread.take() {
+			if let Err(_) = thread.join() {
+				error!(target: LOG_TARGET, "Errored stopping the thread");
+			}
+		}
+	}
+}
+
+async fn listen_for_chunks<PGM, S>(
+	p: PGM,
+	topic: Hash,
+	mut sender: S
+)
+where
+	PGM: ProvideGossipMessages,
+	S: Sink<WorkerMsg> + Unpin,
+{
+	trace!(target: LOG_TARGET, "Registering gossip listener for topic {}", topic);
+	let mut chunks_stream = p.gossip_messages_for(topic);
+
+	while let Some(item) = chunks_stream.next().await {
+		let (s, _) = oneshot::channel();
+		trace!(target: LOG_TARGET, "Received for {:?}", item);
+		let chunks = Chunks {
+			relay_parent: item.0,
+			candidate_hash: item.1,
+			chunks: vec![item.2],
+			result: s,
+		};
+
+		if let Err(_) = sender.send(WorkerMsg::Chunks(chunks)).await {
+			break;
+		}
+	}
+}
+
+
+fn fetch_candidates<P>(client: &P, block: &BlockId, parent: &BlockId)
+	-> ClientResult<Option<impl Iterator<Item=CandidateReceipt>>>
+where
+	P: BlockBody<Block> + ProvideRuntimeApi,
+	P::Api: ParachainHost<Block> + ApiExt<Block, Error=sp_blockchain::Error>,
+{
+	let extrinsics = client.block_body(block)?;
+	Ok(match extrinsics {
+		Some(extrinsics) => client.runtime_api()
+			.get_heads(&parent, extrinsics).map_err(|_| ConsensusError::ChainLookup("".into()))?
+			.and_then(|v| Some(v.into_iter())),
+		None => None,
+	})
+}
+
+/// Creates a task to prune entries in availability store upon block finalization.
+async fn prune_unneeded_availability<P, S>(client: Arc<P>, mut sender: S)
+where
+	P: ProvideRuntimeApi + BlockchainEvents<Block> + BlockBody<Block> + Send + Sync + 'static,
+	P::Api: ParachainHost<Block> + ApiExt<Block, Error=sp_blockchain::Error>,
+	S: Sink<WorkerMsg> + Clone + Send + Sync + Unpin,
+{
+	let mut finality_notification_stream = client.finality_notification_stream();
+
+	while let Some(notification) = finality_notification_stream.next().await {
+		let hash = notification.hash;
+		let parent_hash = notification.header.parent_hash;
+		let candidate_hashes = match fetch_candidates(
+			&*client,
+			&BlockId::hash(hash),
+			&BlockId::hash(parent_hash)
+		) {
+			Ok(Some(candidates)) => candidates.map(|c| c.hash()).collect(),
+			Ok(None) => {
+				warn!(
+					target: LOG_TARGET,
+					"Failed to extract candidates from block body of imported block {:?}", hash
+				);
+				continue;
+			}
+			Err(e) => {
+				warn!(
+					target: LOG_TARGET,
+					"Failed to fetch block body for imported block {:?}: {:?}", hash, e
+				);
+				continue;
+			}
+		};
+
+		let msg = WorkerMsg::CandidatesFinalized(CandidatesFinalized {
+			relay_parent: parent_hash,
+			candidate_hashes
+		});
+
+		if let Err(_) = sender.send(msg).await {
+			break;
+		}
+	}
+}
+
+impl<PGM> Drop for Worker<PGM> {
+	fn drop(&mut self) {
+		for (_, signal) in self.registered_gossip_streams.drain() {
+			signal.fire();
+		}
+	}
+}
+
+impl<PGM> Worker<PGM>
+where
+	PGM: ProvideGossipMessages + Clone + Send + 'static,
+{
+
+	// Called on startup of the worker to register listeners for all awaited chunks.
+	fn register_listeners(
+		&mut self,
+		runtime_handle: &mut Handle,
+		sender: &mut mpsc::UnboundedSender<WorkerMsg>,
+	) {
+		if let Some(awaited_chunks) = self.availability_store.awaited_chunks() {
+			for chunk in awaited_chunks {
+				if let Err(e) = self.register_chunks_listener(
+					runtime_handle,
+					sender,
+					chunk.0,
+					chunk.1,
+				) {
+					warn!(target: LOG_TARGET, "Failed to register gossip listener: {}", e);
+				}
+			}
+		}
+	}
+
+	fn register_chunks_listener(
+		&mut self,
+		runtime_handle: &mut Handle,
+		sender: &mut mpsc::UnboundedSender<WorkerMsg>,
+		relay_parent: Hash,
+		erasure_root: Hash,
+	) -> Result<(), Error> {
+		let (local_id, _) = self.availability_store
+			.get_validator_index_and_n_validators(&relay_parent)
+			.ok_or(Error::IdAndNValidatorsNotFound { relay_parent })?;
+		let topic = erasure_coding_topic(relay_parent, erasure_root, local_id);
+		trace!(
+			target: LOG_TARGET,
+			"Registering listener for erasure chunks topic {} for ({}, {})",
+			topic,
+			relay_parent,
+			erasure_root,
+		);
+
+		let (signal, exit) = exit_future::signal();
+
+		let fut = listen_for_chunks(
+			self.provide_gossip_messages.clone(),
+			topic,
+			sender.clone(),
+		);
+
+		self.registered_gossip_streams.insert(topic, signal);
+
+		let _ = runtime_handle.spawn(
+			fut
+				.unit_error()
+				.boxed()
+				.compat()
+				.select(exit)
+				.then(|_| Ok(()))
+			);
+
+		Ok(())
+	}
+
+	fn on_parachain_blocks_received(
+		&mut self,
+		runtime_handle: &mut Handle,
+		sender: &mut mpsc::UnboundedSender<WorkerMsg>,
+		relay_parent: Hash,
+		blocks: Vec<(CandidateReceipt, Option<(BlockData, AvailableMessages)>)>,
+	) -> Result<(), Error> {
+		let hashes: Vec<_> = blocks.iter().map(|(c, _)| c.hash()).collect();
+
+		// First we have to add the receipts themselves.
+		for (candidate, block) in blocks.into_iter() {
+			let _ = self.availability_store.add_candidate(&candidate);
+
+			if let Some((_block, _msgs)) = block {
+				// Should we be breaking block into chunks here and gossiping it and so on?
+			}
+
+			if let Err(e) = self.register_chunks_listener(
+				runtime_handle,
+				sender,
+				relay_parent,
+				candidate.erasure_root
+			) {
+				warn!(target: LOG_TARGET, "Failed to register chunk listener: {}", e);
+			}
+		}
+
+		let _ = self.availability_store.add_candidates_in_relay_block(
+			&relay_parent,
+			hashes
+		);
+
+		Ok(())
+	}
+
+	// Processes chunks messages that contain awaited items.
+	//
+	// When an awaited item is received, it is placed into the availability store
+	// and removed from the frontier. Listener de-registered.
+	fn on_chunks_received(
+		&mut self,
+		relay_parent: Hash,
+		candidate_hash: Hash,
+		chunks: Vec<ErasureChunk>,
+	) -> Result<(), Error> {
+		let (_, n_validators) = self.availability_store
+			.get_validator_index_and_n_validators(&relay_parent)
+			.ok_or(Error::IdAndNValidatorsNotFound { relay_parent })?;
+
+		let receipt = self.availability_store.get_candidate(&candidate_hash)
+			.ok_or(Error::CandidateNotFound { candidate_hash })?;
+
+		for chunk in &chunks {
+			let topic = erasure_coding_topic(relay_parent, receipt.erasure_root, chunk.index);
+			// need to remove gossip listener and stop it.
+			if let Some(signal) = self.registered_gossip_streams.remove(&topic) {
+				signal.fire();
+			}
+		}
+
+		self.availability_store.add_erasure_chunks(
+			n_validators,
+			&relay_parent,
+			&candidate_hash,
+			chunks,
+		)?;
+
+		Ok(())
+	}
+
+	// Adds the erasure roots into the store.
+	fn on_erasure_roots_received(
+		&mut self,
+		relay_parent: Hash,
+		erasure_roots: Vec<Hash>
+	) -> Result<(), Error> {
+		self.availability_store.add_erasure_roots_in_relay_block(&relay_parent, erasure_roots)?;
+
+		Ok(())
+	}
+
+	// Processes the `ListenForChunks` message.
+	//
+	// When the worker receives a `ListenForChunk` message, it double-checks that
+	// we don't have that piece, and then it registers a listener.
+	fn on_listen_for_chunks_received(
+		&mut self,
+		runtime_handle: &mut Handle,
+		sender: &mut mpsc::UnboundedSender<WorkerMsg>,
+		relay_parent: Hash,
+		candidate_hash: Hash,
+		id: usize
+	) -> Result<(), Error> {
+		let candidate = self.availability_store.get_candidate(&candidate_hash)
+			.ok_or(Error::CandidateNotFound { candidate_hash })?;
+
+		if self.availability_store
+			.get_erasure_chunk(&relay_parent, candidate.block_data_hash, id)
+			.is_none() {
+			if let Err(e) = self.register_chunks_listener(
+				runtime_handle,
+				sender,
+				relay_parent,
+				candidate.erasure_root
+			) {
+				warn!(target: LOG_TARGET, "Failed to register a gossip listener: {}", e);
+			}
+		}
+
+		Ok(())
+	}
+
+	/// Starts a worker with a given availability store and a gossip messages provider.
+	pub fn start(
+		availability_store: Store,
+		provide_gossip_messages: PGM,
+	) -> WorkerHandle {
+		let (sender, mut receiver) = mpsc::unbounded();
+
+		let mut worker = Self {
+			availability_store,
+			provide_gossip_messages,
+			registered_gossip_streams: HashMap::new(),
+			sender: sender.clone(),
+		};
+
+		let sender = sender.clone();
+		let (signal, exit) = exit_future::signal();
+
+		let handle = thread::spawn(move || -> io::Result<()> {
+			let mut runtime = LocalRuntime::new()?;
+			let mut sender = worker.sender.clone();
+
+			let mut runtime_handle = runtime.handle();
+
+			// On startup, registers listeners (gossip streams) for all
+			// (relay_parent, erasure-root, i) in the awaited frontier.
+			worker.register_listeners(&mut runtime_handle, &mut sender);
+
+			let process_notification = async move {
+				while let Some(msg) = receiver.next().await {
+					trace!(target: LOG_TARGET, "Received message {:?}", msg);
+
+					let res = match msg {
+						WorkerMsg::ErasureRoots(msg) => {
+							let ErasureRoots { relay_parent, erasure_roots, result} = msg;
+							let res = worker.on_erasure_roots_received(
+								relay_parent,
+								erasure_roots,
+							);
+							let _ = result.send(res);
+							Ok(())
+						}
+						WorkerMsg::ListenForChunks(msg) => {
+							let ListenForChunks {
+								relay_parent,
+								candidate_hash,
+								index,
+								result,
+							} = msg;
+
+							let res = worker.on_listen_for_chunks_received(
+								&mut runtime_handle,
+								&mut sender,
+								relay_parent,
+								candidate_hash,
+								index as usize,
+							);
+
+							if let Some(result) = result {
+								let _ = result.send(res);
+							}
+							Ok(())
+						}
+						WorkerMsg::ParachainBlocks(msg) => {
+							let ParachainBlocks {
+								relay_parent,
+								blocks,
+								result,
+							} = msg;
+
+							let res = worker.on_parachain_blocks_received(
+								&mut runtime_handle,
+								&mut sender,
+								relay_parent,
+								blocks,
+							);
+
+							let _ = result.send(res);
+							Ok(())
+						}
+						WorkerMsg::Chunks(msg) => {
+							let Chunks { relay_parent, candidate_hash, chunks, result } = msg;
+							let res = worker.on_chunks_received(
+								relay_parent,
+								candidate_hash,
+								chunks,
+							);
+
+							let _ = result.send(res);
+							Ok(())
+						}
+						WorkerMsg::CandidatesFinalized(msg) => {
+							let CandidatesFinalized { relay_parent, candidate_hashes } = msg;
+
+							worker.availability_store.candidates_finalized(
+								relay_parent,
+								candidate_hashes.into_iter().collect(),
+							)
+						}
+						WorkerMsg::MakeAvailable(msg) => {
+							let MakeAvailable { data, result } = msg;
+							let res = worker.availability_store.make_available(data)
+								.map_err(|e| e.into());
+							let _ = result.send(res);
+							Ok(())
+						}
+					};
+
+					if let Err(_) = res {
+						warn!(target: LOG_TARGET, "An error occured while processing a message");
+					}
+				}
+
+			};
+
+			runtime.spawn(
+				process_notification
+					.unit_error()
+					.boxed()
+					.compat()
+					.select(exit.clone())
+					.then(|_| Ok(()))
+			);
+
+			if let Err(e) = runtime.block_on(exit) {
+				warn!(target: LOG_TARGET, "Availability worker error {:?}", e);
+			}
+
+			info!(target: LOG_TARGET, "Availability worker exiting");
+
+			Ok(())
+		});
+
+		WorkerHandle {
+			thread: Some(handle),
+			sender,
+			exit_signal: Some(signal),
+		}
+	}
+}
+
+/// Implementor of the [`BlockImport`] trait.
+///
+/// Used to embed `availability-store` logic into the block imporing pipeline.
+///
+/// [`BlockImport`]: https://substrate.dev/rustdocs/v1.0/substrate_consensus_common/trait.BlockImport.html
+pub struct AvailabilityBlockImport<I, P> {
+	availability_store: Store,
+	inner: I,
+	client: Arc<P>,
+	keystore: KeyStorePtr,
+	to_worker: mpsc::UnboundedSender<WorkerMsg>,
+	exit_signal: Option<exit_future::Signal>,
+}
+
+impl<I, P> Drop for AvailabilityBlockImport<I, P> {
+	fn drop(&mut self) {
+		if let Some(signal) = self.exit_signal.take() {
+			signal.fire();
+		}
+	}
+}
+
+impl<I, P> BlockImport<Block> for AvailabilityBlockImport<I, P> where
+	I: BlockImport<Block> + Send + Sync,
+	I::Error: Into<ConsensusError>,
+	P: ProvideRuntimeApi + ProvideCache<Block>,
+	P::Api: ParachainHost<Block>,
+	P::Api: ApiExt<Block, Error = sp_blockchain::Error>,
+{
+	type Error = ConsensusError;
+
+	fn import_block(
+		&mut self,
+		block: BlockImportParams<Block>,
+		new_cache: HashMap<CacheKeyId, Vec<u8>>,
+	) -> Result<ImportResult, Self::Error> {
+		trace!(
+			target: LOG_TARGET,
+			"Importing block #{}, ({})",
+			block.header.number(),
+			block.post_header().hash()
+		);
+
+		if let Some(ref extrinsics) = block.body {
+			let relay_parent = *block.header.parent_hash();
+			let parent_id = BlockId::hash(*block.header.parent_hash());
+			// Extract our local position i from the validator set of the parent.
+			let validators = self.client.runtime_api().validators(&parent_id)
+				.map_err(|e| ConsensusError::ChainLookup(e.to_string()))?;
+
+			let our_id = self.our_id(&validators);
+
+			// Use a runtime API to extract all included erasure-roots from the imported block.
+			let candidates = self.client.runtime_api().get_heads(&parent_id, extrinsics.clone())
+				.map_err(|e| ConsensusError::ChainLookup(e.to_string()))?;
+
+			match candidates {
+				Some(candidates) => {
+					match our_id {
+						Some(our_id) => {
+							trace!(
+								target: LOG_TARGET,
+								"Our validator id is {}, the candidates included are {:?}",
+								our_id, candidates
+							);
+
+							for candidate in &candidates {
+								// If we don't yet have our chunk of this candidate,
+								// tell the worker to listen for one.
+								if self.availability_store.get_erasure_chunk(
+									&relay_parent,
+									candidate.block_data_hash,
+									our_id as usize,
+								).is_none() {
+									let msg = WorkerMsg::ListenForChunks(ListenForChunks {
+										relay_parent,
+										candidate_hash: candidate.hash(),
+										index: our_id as u32,
+										result: None,
+									});
+
+									let _ = self.to_worker.unbounded_send(msg);
+								}
+							}
+
+							let erasure_roots: Vec<_> = candidates
+								.iter()
+								.map(|c| c.erasure_root)
+								.collect();
+
+							// Inform the worker about new (relay_parent, erasure_roots) pairs
+							let (s, _) = oneshot::channel();
+							let msg = WorkerMsg::ErasureRoots(ErasureRoots {
+								relay_parent,
+								erasure_roots,
+								result: s,
+							});
+
+							let _ = self.to_worker.unbounded_send(msg);
+
+							let (s, _) = oneshot::channel();
+
+							// Inform the worker about the included parachain blocks.
+							let msg = WorkerMsg::ParachainBlocks(ParachainBlocks {
+								relay_parent,
+								blocks: candidates.into_iter().map(|c| (c, None)).collect(),
+								result: s,
+							});
+
+							let _ = self.to_worker.unbounded_send(msg);
+						}
+						None => (),
+					}
+				}
+				None => {
+					trace!(
+						target: LOG_TARGET,
+						"No parachain heads were included in block {}", block.header.hash()
+					);
+				},
+			}
+		}
+
+		self.inner.import_block(block, new_cache).map_err(Into::into)
+	}
+
+	fn check_block(
+		&mut self,
+		block: BlockCheckParams<Block>,
+	) -> Result<ImportResult, Self::Error> {
+		self.inner.check_block(block).map_err(Into::into)
+	}
+}
+
+impl<I, P> AvailabilityBlockImport<I, P> {
+	pub(crate) fn new(
+		availability_store: Store,
+		client: Arc<P>,
+		block_import: I,
+		thread_pool: TaskExecutor,
+		keystore: KeyStorePtr,
+		to_worker: mpsc::UnboundedSender<WorkerMsg>,
+	) -> Self
+	where
+		P: ProvideRuntimeApi + BlockBody<Block> + BlockchainEvents<Block> + Send + Sync + 'static,
+		P::Api: ParachainHost<Block>,
+		P::Api: ApiExt<Block, Error = sp_blockchain::Error>,
+	{
+		let (signal, exit) = exit_future::signal();
+
+		// This is not the right place to spawn the finality future,
+		// it would be more appropriate to spawn it in the `start` method of the `Worker`.
+		// However, this would make the type of the `Worker` and the `Store` itself
+		// dependent on the types of client and executor, which would prove
+		// not not so handy in the testing code.
+		let mut exit_signal = Some(signal);
+		let prune_available = prune_unneeded_availability(client.clone(), to_worker.clone())
+			.unit_error()
+			.boxed()
+			.compat()
+			.select(exit.clone())
+			.then(|_| Ok(()));
+
+		if let Err(_) = thread_pool.execute(Box::new(prune_available)) {
+			error!(target: LOG_TARGET, "Failed to spawn availability pruning task");
+			exit_signal = None;
+		}
+
+		AvailabilityBlockImport {
+			availability_store,
+			client,
+			inner: block_import,
+			to_worker,
+			keystore,
+			exit_signal,
+		}
+	}
+
+	fn our_id(&self, validators: &[ValidatorId]) -> Option<u32> {
+		let keystore = self.keystore.read();
+		validators
+			.iter()
+			.enumerate()
+			.find_map(|(i, v)| {
+				keystore.key_pair::<ValidatorPair>(&v).map(|_| i as u32).ok()
+			})
+	}
+}
+
+
+#[cfg(test)]
+mod tests {
+	use super::*;
+	use std::time::Duration;
+	use futures::{stream, channel::mpsc, Stream};
+	use std::sync::{Arc, Mutex};
+	use tokio::runtime::Runtime;
+
+	// Just contains topic->channel mapping to give to outer code on `gossip_messages_for` calls.
+	struct TestGossipMessages {
+		messages: Arc<Mutex<HashMap<Hash, mpsc::UnboundedReceiver<(Hash, Hash, ErasureChunk)>>>>,
+	}
+
+	impl ProvideGossipMessages for TestGossipMessages {
+		fn gossip_messages_for(&self, topic: Hash)
+			-> Box<dyn Stream<Item = (Hash, Hash, ErasureChunk)> + Send + Unpin>
+		{
+			match self.messages.lock().unwrap().remove(&topic) {
+				Some(receiver) => Box::new(receiver),
+				None => Box::new(stream::iter(vec![])),
+			}
+		}
+
+		fn gossip_erasure_chunk(
+			&self,
+			_relay_parent: Hash,
+			_candidate_hash: Hash,
+			_erasure_root: Hash,
+			_chunk: ErasureChunk
+		) {}
+	}
+
+	impl Clone for TestGossipMessages {
+		fn clone(&self) -> Self {
+			TestGossipMessages {
+				messages: self.messages.clone(),
+			}
+		}
+	}
+
+	// This test tests that as soon as the worker receives info about new parachain blocks
+	// included it registers gossip listeners for it's own chunks. Upon receiving the awaited
+	// chunk messages the corresponding listeners are deregistered and these chunks are removed
+	// from the awaited chunks set.
+	#[test]
+	fn receiving_gossip_chunk_removes_from_frontier() {
+		let mut runtime = Runtime::new().unwrap();
+		let relay_parent = [1; 32].into();
+		let erasure_root = [2; 32].into();
+		let local_id = 2;
+		let n_validators = 4;
+
+		let store = Store::new_in_memory();
+
+		// Tell the store our validator's position and the number of validators at given point.
+		store.add_validator_index_and_n_validators(&relay_parent, local_id, n_validators).unwrap();
+
+		let (gossip_sender, gossip_receiver) = mpsc::unbounded();
+
+		let topic = erasure_coding_topic(relay_parent, erasure_root, local_id);
+
+		let messages = TestGossipMessages {
+			messages: Arc::new(Mutex::new(vec![
+						  (topic, gossip_receiver)
+			].into_iter().collect()))
+		};
+
+		let mut candidate = CandidateReceipt::default();
+
+		candidate.erasure_root = erasure_root;
+		let candidate_hash = candidate.hash();
+
+		// At this point we shouldn't be waiting for any chunks.
+		assert!(store.awaited_chunks().is_none());
+
+		let (s, r) = oneshot::channel();
+
+		let msg = WorkerMsg::ParachainBlocks(ParachainBlocks {
+			relay_parent,
+			blocks: vec![(candidate, None)],
+			result: s,
+		});
+
+		let handle = Worker::start(store.clone(), messages);
+
+		// Tell the worker that the new blocks have been included into the relay chain.
+		// This should trigger the registration of gossip message listeners for the
+		// chunk topics.
+		handle.sender.unbounded_send(msg).unwrap();
+
+		runtime.block_on(r.unit_error().boxed().compat()).unwrap().unwrap().unwrap();
+
+		// Make sure that at this point we are waiting for the appropriate chunk.
+		assert_eq!(
+			store.awaited_chunks().unwrap(),
+			vec![(relay_parent, erasure_root, candidate_hash, local_id)].into_iter().collect()
+		);
+
+		let msg = (
+			relay_parent,
+			candidate_hash,
+			ErasureChunk {
+				chunk: vec![1, 2, 3],
+				index: local_id as u32,
+				proof: vec![],
+			}
+		);
+
+		// Send a gossip message with an awaited chunk
+		gossip_sender.unbounded_send(msg).unwrap();
+
+		// At the point the needed piece is received, the gossip listener for
+		// this topic is deregistered and it's receiver side is dropped.
+		// Wait for the sender side to become closed.
+		while !gossip_sender.is_closed() {
+			// Probably we can just .wait this somehow?
+			thread::sleep(Duration::from_millis(100));
+		}
+
+		// The awaited chunk has been received so at this point we no longer wait for any chunks.
+		assert_eq!(store.awaited_chunks().unwrap().len(), 0);
+	}
+
+	#[test]
+	fn listen_for_chunk_registers_listener() {
+		let mut runtime = Runtime::new().unwrap();
+		let relay_parent = [1; 32].into();
+		let erasure_root_1 = [2; 32].into();
+		let erasure_root_2 = [3; 32].into();
+		let block_data_hash_1 = [4; 32].into();
+		let block_data_hash_2 = [5; 32].into();
+		let local_id = 2;
+		let n_validators = 4;
+
+		let mut candidate_1 = CandidateReceipt::default();
+		candidate_1.erasure_root = erasure_root_1;
+		candidate_1.block_data_hash = block_data_hash_1;
+		let candidate_1_hash = candidate_1.hash();
+
+		let mut candidate_2 = CandidateReceipt::default();
+		candidate_2.erasure_root = erasure_root_2;
+		candidate_2.block_data_hash = block_data_hash_2;
+		let candidate_2_hash = candidate_2.hash();
+
+		let store = Store::new_in_memory();
+
+		// Tell the store our validator's position and the number of validators at given point.
+		store.add_validator_index_and_n_validators(&relay_parent, local_id, n_validators).unwrap();
+
+		// Let the store know about the candidates
+		store.add_candidate(&candidate_1).unwrap();
+		store.add_candidate(&candidate_2).unwrap();
+
+		// And let the store know about the chunk from the second candidate.
+		store.add_erasure_chunks(
+			n_validators,
+			&relay_parent,
+			&candidate_2_hash,
+			vec![ErasureChunk {
+				chunk: vec![1, 2, 3],
+				index: local_id,
+				proof: Vec::default(),
+			}],
+		).unwrap();
+
+		let (_, gossip_receiver_1) = mpsc::unbounded();
+		let (_, gossip_receiver_2) = mpsc::unbounded();
+
+		let topic_1 = erasure_coding_topic(relay_parent, erasure_root_1, local_id);
+		let topic_2 = erasure_coding_topic(relay_parent, erasure_root_2, local_id);
+
+		let messages = TestGossipMessages {
+			messages: Arc::new(Mutex::new(
+			vec![
+				(topic_1, gossip_receiver_1),
+				(topic_2, gossip_receiver_2),
+			].into_iter().collect()))
+		};
+
+		let handle = Worker::start(store.clone(), messages.clone());
+
+		let (s2, r2) = oneshot::channel();
+		// Tell the worker to listen for chunks from candidate 2 (we alredy have a chunk from it).
+		let listen_msg_2 = WorkerMsg::ListenForChunks(ListenForChunks {
+			relay_parent,
+			candidate_hash: candidate_2_hash,
+			index: local_id as u32,
+			result: Some(s2),
+		});
+
+		handle.sender.unbounded_send(listen_msg_2).unwrap();
+
+		runtime.block_on(r2.unit_error().boxed().compat()).unwrap().unwrap().unwrap();
+		// The gossip sender for this topic left intact => listener not registered.
+		assert!(messages.messages.lock().unwrap().contains_key(&topic_2));
+
+		let (s1, r1) = oneshot::channel();
+
+		// Tell the worker to listen for chunks from candidate 1.
+		// (we don't have a chunk from it yet).
+		let listen_msg_1 = WorkerMsg::ListenForChunks(ListenForChunks {
+			relay_parent,
+			candidate_hash: candidate_1_hash,
+			index: local_id as u32,
+			result: Some(s1),
+		});
+
+		handle.sender.unbounded_send(listen_msg_1).unwrap();
+		runtime.block_on(r1.unit_error().boxed().compat()).unwrap().unwrap().unwrap();
+
+		// The gossip sender taken => listener registered.
+		assert!(!messages.messages.lock().unwrap().contains_key(&topic_1));
+	}
+}