feat: Detach chunk identity from chunk name

Cargo.toml

@@ -7,7 +7,7 @@ edition = "2021"
 [features]
 query-tracing = ["sqd-query/max_level_trace"]
 
-[dependencies] 
+[dependencies]
 anyhow = "1.0"
 async-compression = { version = "0.4.27", features = ["gzip", "tokio"] }
 async-stream = "0.3.5"
@@ -59,17 +59,20 @@ tracing-subscriber = { version = "0.3.18", features = ["env-filter"] }
 url = "2.5.2"
 walkdir = "2.5.0"
 zstd = "0.13"
- 
-sqd-assignments = { git = "https://github.com/subsquid/sqd-network.git", rev = "a1b685d", features = ["reader"] }
-sqd-contract-client = { git = "https://github.com/subsquid/sqd-network.git", rev = "a1b685d", version = "1.2.1" }
-sqd-messages = { git = "https://github.com/subsquid/sqd-network.git", rev = "a1b685d", version = "2.0.2", features = ["bitstring"] }
-sqd-network-transport = { git = "https://github.com/subsquid/sqd-network.git", rev = "a1b685d", version = "3.0.0", features = ["worker", "metrics"] }
+
+sqd-assignments = { git = "https://github.com/subsquid/sqd-network.git", rev = "938ffe1", features = ["reader"] }
+sqd-contract-client = { git = "https://github.com/subsquid/sqd-network.git", rev = "938ffe1", version = "1.2.1" }
+sqd-messages = { git = "https://github.com/subsquid/sqd-network.git", rev = "938ffe1", version = "2.0.2", features = ["bitstring"] }
+sqd-network-transport = { git = "https://github.com/subsquid/sqd-network.git", rev = "938ffe1", version = "3.0.0", features = ["worker", "metrics"] }
 
 sqd-query = { git = "https://github.com/subsquid/data.git", rev = "9db54d3", features = ["parquet"] }
 sqd-polars = { git = "https://github.com/subsquid/data.git", rev = "9db54d3" }
 
 sql_query_plan = {git = "https://github.com/subsquid/qplan.git", rev = "658f88f" }
 
+[dev-dependencies]
+sqd-assignments = { git = "https://github.com/subsquid/sqd-network.git", rev = "938ffe1", features = ["builder"] }
+
 [profile.release]
 debug = true
 opt-level = 3

src/controller/p2p.rs

@@ -465,18 +465,27 @@ impl<EventStream: Stream<Item = WorkerEvent> + Send + 'static> P2PController<Eve
             }
         }
 
+        let Some(block_range) = query
+            .block_range
+            .map(|sqd_messages::Range { begin, end }| (begin, end))
+        else {
+            return (
+                Err(QueryError::BadRequest("block_range is required".to_owned())),
+                None,
+            );
+        };
+
         let mut allocation_chip = 1.0f32;
 
         if let Ok(chunk) = query.chunk_id.parse::<DataChunk>() {
-            if let Some(range) = query.block_range {
-                let active_len = std::cmp::min(chunk.last_block.into(), range.end)
-                    .saturating_sub(std::cmp::max(chunk.first_block.into(), range.begin))
-                    .max(1);
-                let chunk_len = Into::<u64>::into(chunk.last_block)
-                    .saturating_sub(chunk.first_block.into())
-                    .max(1);
-                allocation_chip = active_len as f32 / chunk_len as f32;
-            }
+            let (begin, end) = block_range;
+            let active_len = std::cmp::min(chunk.last_block.into(), end)
+                .saturating_sub(std::cmp::max(chunk.first_block.into(), begin))
+                .max(1);
+            let chunk_len = Into::<u64>::into(chunk.last_block)
+                .saturating_sub(chunk.first_block.into())
+                .max(1);
+            allocation_chip = active_len as f32 / chunk_len as f32;
         };
 
         // We claim 1. allocation first and refund unused allocation later. It's done to prevent burst overloading with small requests.
@@ -493,9 +502,6 @@ impl<EventStream: Stream<Item = WorkerEvent> + Send + 'static> P2PController<Eve
         };
         let mut retry_after = status.retry_after();
 
-        let block_range = query
-            .block_range
-            .map(|sqd_messages::Range { begin, end }| (begin, end));
         let result = self
             .worker
             .run_query(

src/controller/worker.rs

@@ -22,10 +22,7 @@ use crate::{
     controller::{polars_target, sql_request::WorkerChunkStore},
     metrics,
     query::result::{QueryError, QueryOk, QueryResult},
-    storage::{
-        layout::DataChunk,
-        manager::{self, StateManager},
-    },
+    storage::manager::{self, StateManager},
     types::dataset::Dataset,
 };
 
@@ -73,7 +70,7 @@ impl Worker {
         &self,
         query_str: &str,
         dataset: Dataset,
-        block_range: Option<(u64, u64)>,
+        block_range: (u64, u64),
         chunk_id: &str,
         client_id: Option<PeerId>,
         query_type: QueryType,
@@ -112,22 +109,17 @@ impl Worker {
         &self,
         query_str: &str,
         dataset: Dataset,
-        block_range: Option<(u64, u64)>,
+        block_range: (u64, u64),
         chunk_id: &str,
     ) -> QueryResult {
-        let Ok(chunk) = chunk_id.parse::<DataChunk>() else {
-            return Err(QueryError::BadRequest(format!(
-                "Can't parse chunk id '{chunk_id}'"
-            )));
-        };
         let mut query = sqd_query::Query::from_json_bytes(query_str.as_bytes())
             .map_err(|e| QueryError::BadRequest(format!("Couldn't parse query: {e:?}")))?;
-        if let Some((from_block, to_block)) = block_range {
-            query.set_first_block(from_block);
-            query.set_last_block(Some(to_block));
-        }
+        let (from_block, to_block) = block_range;
+
+        query.set_first_block(from_block);
+        query.set_last_block(Some(to_block));
 
-        let Some(chunk_guard) = self.state_manager.clone().get_chunk(dataset, chunk.clone()) else {
+        let Some(chunk_guard) = self.state_manager.clone().get_chunk(dataset, chunk_id) else {
             return Err(QueryError::NotFound);
         };
 
@@ -148,11 +140,12 @@ impl Worker {
                     writer.write_blocks(&mut blocks)?;
                     blocks.last_block()
                 } else {
-                    if let Some(last_query_block) = query.last_block() {
-                        std::cmp::min(last_query_block, chunk.last_block.into())
-                    } else {
-                        chunk.last_block.into()
-                    }
+                    // No matching rows in this chunk. We used to fall back to
+                    // the chunk's last_block (parsed from the chunk_id) so the
+                    // portal could see progress. After NET-385 the worker
+                    // treats chunk_id as opaque, so we report the query's
+                    // upper bound
+                    to_block
                 };
                 let bytes = writer.finish()?;
                 let serialization_duration = serialization_timer.elapsed();
@@ -188,11 +181,6 @@ impl Worker {
         dataset: Dataset,
         chunk_id: &str,
     ) -> QueryResult {
-        let Ok(chunk) = chunk_id.parse::<DataChunk>() else {
-            return Err(QueryError::BadRequest(format!(
-                "Can't parse chunk id '{chunk_id}'"
-            )));
-        };
         let Ok(query_bytes) = base64.decode(query_str) else {
             return Err(QueryError::BadRequest(format!(
                 "Can't decode plan '{query_str}'"
@@ -207,7 +195,7 @@ impl Worker {
         let Some(chunk_guard) = self
             .state_manager
             .clone()
-            .get_chunk(dataset.clone(), chunk.clone())
+            .get_chunk(dataset.clone(), chunk_id)
         else {
             return Err(QueryError::NotFound);
         };

src/storage/datasets_index.rs

@@ -1,22 +1,19 @@
 use std::{collections::HashMap, str::FromStr, sync::Arc};
 
 use reqwest::Url;
-use sqd_network_transport::{Keypair, PeerId};
+use sqd_network_transport::Keypair;
 use tracing::error;
 
-use crate::types::{
-    dataset::Dataset,
-    state::{ChunkRef, ChunkSet},
-};
-
-use super::layout::DataChunk;
+use crate::types::state::ChunkRef;
+use sqd_assignments::ChunkRef as ChunkAssignmentRef;
 
 pub struct DatasetsIndex {
     assignment: sqd_assignments::Assignment,
     assignment_id: String,
-    peer_id: PeerId,
     status: sqd_assignments::WorkerStatus,
     http_headers: reqwest::header::HeaderMap,
+    // chunks assigned to this worker
+    chunks: HashMap<ChunkRef, ChunkAssignmentRef>,
 }
 
 #[derive(Debug, PartialEq, Eq)]
@@ -26,11 +23,12 @@ pub struct RemoteFile {
 }
 
 impl DatasetsIndex {
-    pub fn list_files(&self, dataset: &Dataset, chunk: &DataChunk) -> Option<Vec<RemoteFile>> {
-        let chunk = self
-            .assignment
-            .find_chunk(dataset, *chunk.first_block)
-            .ok()?;
+    /// Returns the remote files (URL + filename) associated with the given
+    /// chunk, or `None` if the chunk is not in the assignment or any URL
+    /// fails to parse.
+    pub fn list_files(&self, chunk: &ChunkRef) -> Option<Vec<RemoteFile>> {
+        let chunk_ref = self.chunks.get(chunk)?;
+        let chunk = self.assignment.get_chunk(*chunk_ref)?;
         let base_url = Url::from_str(&chunk.dataset_base_url())
             .inspect_err(|e| {
                 tracing::warn!(
@@ -81,42 +79,25 @@ impl DatasetsIndex {
             })
             .collect();
 
+        let mut chunks = HashMap::new();
+        let mut pool = StringPool::default();
+        for (chunk_ref, chunk) in worker.iter_chunks_with_ref() {
+            let key = ChunkRef {
+                dataset: pool.get(chunk.dataset_id()),
+                chunk: Arc::from(chunk.id()),
+            };
+            chunks.insert(key, chunk_ref);
+        }
+
         Ok(Self {
             status: worker.status(),
             assignment,
             assignment_id: id.into(),
-            peer_id,
             http_headers,
+            chunks,
         })
     }
 
-    pub fn create_chunks_set(&self) -> ChunkSet {
-        let mut chunk_set = ChunkSet::new();
-        let Some(worker) = self.assignment.get_worker(&self.peer_id) else {
-            return chunk_set;
-        };
-        let mut pool = StringPool::default();
-        for chunk in worker.iter_chunks() {
-            match DataChunk::from_str(chunk.id()) {
-                Ok(id) => {
-                    let chunk = ChunkRef {
-                        dataset: pool.get(chunk.dataset_id()),
-                        chunk: id,
-                    };
-                    if let Some(last) = chunk_set.last() {
-                        debug_assert!(
-                            last < &chunk,
-                            "Assigned chunks are not sorted: {last} >= {chunk}"
-                        );
-                    }
-                    chunk_set.insert(chunk);
-                }
-                Err(e) => tracing::warn!("Couldn't parse chunk id {}: {e}", chunk.id()),
-            }
-        }
-        chunk_set
-    }
-
     pub fn status(&self) -> sqd_assignments::WorkerStatus {
         self.status
     }
@@ -128,6 +109,10 @@ impl DatasetsIndex {
     pub fn assignment_id(&self) -> &str {
         &self.assignment_id
     }
+
+    pub fn chunks(&self) -> &HashMap<ChunkRef, ChunkAssignmentRef> {
+        &self.chunks
+    }
 }
 
 #[derive(Default)]

src/storage/downloader.rs

@@ -11,11 +11,7 @@ use tracing::instrument;
 
 use crate::{cli, types::state::ChunkRef};
 
-use super::{
-    datasets_index::{DatasetsIndex, RemoteFile},
-    guard::FsGuard,
-    local_fs::add_temp_prefix,
-};
+use super::{datasets_index::RemoteFile, guard::FsGuard, local_fs::add_temp_prefix};
 
 const START_DELAY: Duration = Duration::from_millis(100);
 
@@ -51,7 +47,8 @@ impl ChunkDownloader {
         &mut self,
         chunk: ChunkRef,
         dst: PathBuf,
-        datasets_index: &DatasetsIndex,
+        files: Vec<RemoteFile>,
+        headers: reqwest::header::HeaderMap,
     ) {
         let cancel_token = CancellationToken::new();
 
@@ -62,13 +59,7 @@ impl ChunkDownloader {
             panic!("Chunk {chunk} is already being downloaded");
         }
 
-        let files = datasets_index
-            .list_files(&chunk.dataset, &chunk.chunk)
-            .unwrap_or_else(|| {
-                panic!("Dataset {} not found", chunk.dataset);
-            });
         let num_files = files.len();
-        let headers = datasets_index.get_headers().clone();
         let client = self.reqwest_client.clone();
         let current_delay = self.current_delay;
         let s3_timeout = self.args.s3_timeout;