MultiForkByKeyProvider for all datagen

provider/blob/src/export/blob_exporter.rs

@@ -7,31 +7,34 @@ use icu_provider::export::DataExporter;
 use icu_provider::iter::IterableDynProvider;
 use icu_provider::prelude::*;
 use icu_provider::serde::SerializeMarker;
-use litemap::LiteMap;
 use writeable::Writeable;
 use zerovec::map2d::ZeroMap2d;
+use std::sync::Mutex;
 
 /// A data exporter that writes data to a single-file blob.
 /// See the module-level docs for an example.
 pub struct BlobExporter {
-    resources: LiteMap<(ResourceKeyHash, String), Vec<u8>>,
-    sink: Box<dyn std::io::Write>,
+    resources: Mutex<Vec<(ResourceKeyHash, String, Vec<u8>)>>,
+    sink: Mutex<Option<Box<dyn std::io::Write>>>,
 }
 
+unsafe impl Sync for BlobExporter {}
+unsafe impl Send for BlobExporter {}
+
 impl BlobExporter {
     /// Create a [`BlobExporter`] that writes to the given I/O stream.
     pub fn new_with_sink(sink: Box<dyn std::io::Write>) -> Self {
         Self {
-            resources: LiteMap::new(),
-            sink,
+            resources: Mutex::new(Vec::new()),
+            sink: Mutex::new(Some(sink)),
         }
     }
 }
 
 impl Drop for BlobExporter {
     fn drop(&mut self) {
-        if !self.resources.is_empty() {
-            panic!("Please call close before dropping FilesystemExporter");
+        if self.sink.lock().unwrap().is_some() {
+            panic!("Please call close before dropping BlobExporter");
         }
     }
 }
@@ -52,34 +55,31 @@ impl DataExporter<SerializeMarker> for BlobExporter {
                 .load_payload(key, req)?
                 .take_payload()?
                 .serialize(&mut <dyn erased_serde::Serializer>::erase(&mut serializer))?;
-            self.resources.insert(
-                (
-                    key.get_hash(),
-                    req.options.writeable_to_string().into_owned(),
-                ),
+            self.resources.lock().unwrap().push((
+                key.get_hash(),
+                req.options.writeable_to_string().into_owned(),
                 serializer.output.0,
-            );
+            ));
         }
         Ok(())
     }
 
     fn close(&mut self) -> Result<(), DataError> {
-        // Convert from LiteMap<(String, String), Vec<u8>> to ZeroMap2d<str, str, [u8]>
-        let mut zm: ZeroMap2d<ResourceKeyHash, str, [u8]> =
-            ZeroMap2d::with_capacity(self.resources.len());
-        for ((key, option), bytes) in self.resources.iter() {
-            zm.insert(key, option, bytes);
+        if let Some(mut sink) = self.sink.lock().unwrap().take() {
+            // We are the first `close` and own the sink now
+            let zm = self.resources.lock().unwrap().drain(..).collect::<ZeroMap2d<_, _, _>>();
+            let blob = BlobSchema::V001(BlobSchemaV1 {
+                resources: zm.as_borrowed(),
+            });
+            log::info!("Serializing blob to output stream...");
+            let mut serializer = postcard::Serializer {
+                output: postcard::flavors::AllocVec(Vec::new()),
+            };
+            serde::Serialize::serialize(&blob, &mut serializer)?;
+            sink.write_all(&serializer.output.0)?;
+            Ok(())
+        } else {
+            Err(DataError::custom("Close was called twice"))
         }
-        let blob = BlobSchema::V001(BlobSchemaV1 {
-            resources: zm.as_borrowed(),
-        });
-        log::info!("Serializing blob to output stream...");
-        let mut serializer = postcard::Serializer {
-            output: postcard::flavors::AllocVec(Vec::new()),
-        };
-        serde::Serialize::serialize(&blob, &mut serializer)?;
-        self.sink.write_all(&serializer.output.0)?;
-        self.resources.clear();
-        Ok(())
     }
 }

provider/fs/src/export/fs_exporter.rs

@@ -59,6 +59,9 @@ pub struct FilesystemExporter {
     serializer: Box<dyn AbstractSerializer>,
 }
 
+unsafe impl Sync for FilesystemExporter {}
+unsafe impl Send for FilesystemExporter {}
+
 impl DataExporter<SerializeMarker> for FilesystemExporter {
     fn put_key_with_options(
         &mut self,