Fix local upload of data objects

Author: spirali

crates/hyperqueue/src/client/commands/data.rs

@@ -25,8 +25,8 @@ pub struct PutOpts {
     /// Path of file/directory that should be uploaded
     path: PathBuf,
     /// DataId of task output
-    #[arg(long, default_value = "")]
-    mime_type: String,
+    #[arg(long)]
+    mime_type: Option<String>,
 }
 
 #[derive(Parser)]

crates/tako/src/internal/datasrv/download.rs

@@ -1,10 +1,10 @@
 use crate::connection::Connection;
 use crate::datasrv::DataObjectId;
+use crate::internal::datasrv::DataObjectRef;
 use crate::internal::datasrv::messages::{
     DataDown, FromDataClientMessage, ToDataClientMessageDown,
 };
 use crate::internal::datasrv::utils::DataObjectComposer;
-use crate::internal::datasrv::DataObjectRef;
 use crate::{Map, WrappedRcRefCell};
 use orion::kex::SecretKey;
 use priority_queue::PriorityQueue;
@@ -13,8 +13,8 @@ use std::rc::Rc;
 use std::sync::Arc;
 use std::time::Duration;
 use tokio::net::TcpStream;
-use tokio::sync::{oneshot, Notify, OwnedSemaphorePermit, Semaphore};
-use tokio::task::{spawn_local, AbortHandle, JoinSet};
+use tokio::sync::{Notify, OwnedSemaphorePermit, Semaphore, oneshot};
+use tokio::task::{AbortHandle, JoinSet, spawn_local};
 use tokio::time::Instant;
 
 const PROTOCOL_VERSION: u32 = 0;

crates/tako/src/internal/datasrv/local_client.rs

@@ -59,11 +59,11 @@ impl LocalDataClient {
     pub async fn put_data_object_from_file(
         &mut self,
         data_id: OutputId,
-        mime_type: String,
+        mime_type: Option<String>,
         path: &Path,
     ) -> crate::Result<()> {
         log::debug!(
-            "Uploading file {} as data_obj={} (mime_type={})",
+            "Uploading file {} as data_obj={} (mime_type={:?})",
             path.display(),
             data_id,
             &mime_type
@@ -81,24 +81,31 @@ impl LocalDataClient {
             if first {
                 self.send_message(FromLocalDataClientMessageUp::PutDataObject {
                     data_id,
-                    mime_type: mime_type.as_str(),
+                    mime_type: mime_type.as_ref().map(|x| x.as_str()),
                     size,
                     data: PutDataUp { data },
                 })
                 .await?;
                 first = false;
-                if bytes_read == 0 {
-                    break;
-                }
             } else {
                 if bytes_read == 0 {
-                    break;
+                    return Err(DsError::GenericError(
+                        "File changed size during upload".to_string(),
+                    ));
                 }
                 self.send_message(FromLocalDataClientMessageUp::PutDataObjectPart(PutDataUp {
                     data,
                 }))
                 .await?;
             }
+            if written as u64 >= size {
+                if written as u64 > size {
+                    return Err(DsError::GenericError(
+                        "File changed size during upload".to_string(),
+                    ));
+                }
+                break;
+            }
         }
         log::debug!("Waiting for confirmation of upload");
         let message = self.read_message().await?;

crates/tako/src/internal/datasrv/messages.rs

@@ -39,7 +39,7 @@ pub(crate) struct PutDataUp<'a> {
 pub(crate) enum FromLocalDataClientMessageUp<'a> {
     PutDataObject {
         data_id: OutputId,
-        mime_type: &'a str,
+        mime_type: Option<&'a str>,
         size: u64,
         data: PutDataUp<'a>,
     },

crates/tako/src/internal/datasrv/tests.rs

@@ -1,8 +1,8 @@
 use crate::datasrv::DataObjectId;
 use crate::internal::datasrv::datastorage::DataStorage;
 use crate::internal::datasrv::test_utils::{
-    start_download_manager, start_test_upload_service, test_download_manager, PlacementConfig,
-    TestUploadInterface,
+    PlacementConfig, TestUploadInterface, start_download_manager, start_test_upload_service,
+    test_download_manager,
 };
 use crate::internal::datasrv::{DataObject, DataObjectRef};
 use crate::internal::tests::utils::sorted_vec;

crates/tako/src/internal/worker/localcomm.rs

@@ -135,6 +135,7 @@ async fn handle_connection(state_ref: WorkerStateRef, stream: UnixStream) -> cra
             .ok_or_else(|| crate::Error::GenericError("Invalid token".to_string()))?;
         match registration {
             Registration::DataConnection { task_id, input_map } => {
+                log::debug!("New local data connection: {task_id}");
                 let task_id = *task_id;
                 let input_map = input_map.clone();
                 drop(lc_state);