Use S3 transfer manager for file uploads

Author: frcroth

app/controllers/DatasetController.scala

@@ -7,7 +7,7 @@ import com.scalableminds.util.objectid.ObjectId
 import com.scalableminds.util.time.Instant
 import com.scalableminds.util.tools.{Fox, TristateOptionJsonHelper}
 import com.scalableminds.webknossos.datastore.models.AdditionalCoordinate
-import com.scalableminds.webknossos.datastore.models.datasource.ElementClass
+import com.scalableminds.webknossos.datastore.models.datasource.{DataSource, ElementClass}
 import mail.{MailchimpClient, MailchimpTag}
 import models.analytics.{AnalyticsService, ChangeDatasetSettingsEvent, OpenDatasetEvent}
 import models.dataset._
@@ -21,7 +21,6 @@ import models.organization.OrganizationDAO
 import models.team.{TeamDAO, TeamService}
 import models.user.{User, UserDAO, UserService}
 import com.scalableminds.util.tools.{Empty, Failure, Full}
-import com.scalableminds.webknossos.datastore.services.DataSourceRegistrationInfo
 import play.api.i18n.{Messages, MessagesProvider}
 import play.api.libs.functional.syntax._
 import play.api.libs.json._
@@ -72,6 +71,12 @@ object SegmentAnythingMaskParameters {
   implicit val jsonFormat: Format[SegmentAnythingMaskParameters] = Json.format[SegmentAnythingMaskParameters]
 }
 
+case class DataSourceRegistrationInfo(dataSource: DataSource, folderId: Option[String], dataStoreName: String)
+
+object DataSourceRegistrationInfo {
+  implicit val jsonFormat: OFormat[DataSourceRegistrationInfo] = Json.format[DataSourceRegistrationInfo]
+}
+
 class DatasetController @Inject()(userService: UserService,
                                   userDAO: UserDAO,
                                   datasetService: DatasetService,

project/Dependencies.scala

@@ -5,6 +5,7 @@ object Dependencies {
   private val silhouetteVersion = "10.0.3"
   private val brotliVersion = "1.19.0"
   private val slickVersion = "3.5.2"
+  private val awsVersion = "2.32.24"
   private val scalapbVersion = scalapb.compiler.Version.scalapbVersion
   private val grpcVersion = scalapb.compiler.Version.grpcJavaVersion
 
@@ -55,7 +56,9 @@ object Dependencies {
     // MultiArray (ndarray) handles. import ucar
     "edu.ucar" % "cdm-core" % "5.4.2",
     // Amazon S3 cloud storage client. import software.amazon.awssdk
-    "software.amazon.awssdk" % "s3" % "2.32.24",
+    "software.amazon.awssdk" % "s3" % awsVersion,
+    // AWS Transfer Manager for multipart uploads. import software.amazon.awssdk.transfer.s3
+    "software.amazon.awssdk" % "s3-transfer-manager" % awsVersion,
     // Google cloud storage client. import com.google.cloud.storage, import com.google.auth.oauth2
     "com.google.cloud" % "google-cloud-storage" % "2.55.0",
     // Blosc compression. import dev.zarr.bloscjava

webknossos-datastore/app/com/scalableminds/webknossos/datastore/services/DSRemoteWebknossosClient.scala

@@ -50,12 +50,6 @@ object MagPathInfo {
   implicit val jsonFormat: OFormat[MagPathInfo] = Json.format[MagPathInfo]
 }
 
-case class DataSourceRegistrationInfo(dataSource: DataSource, folderId: Option[String], dataStoreName: String)
-
-object DataSourceRegistrationInfo {
-  implicit val jsonFormat: OFormat[DataSourceRegistrationInfo] = Json.format[DataSourceRegistrationInfo]
-}
-
 trait RemoteWebknossosClient {
   def requestUserAccess(accessRequest: UserAccessRequest)(implicit tc: TokenContext): Fox[UserAccessAnswer]
 }

webknossos-datastore/app/com/scalableminds/webknossos/datastore/services/uploading/UploadService.scala

@@ -31,7 +31,6 @@ import com.typesafe.scalalogging.LazyLogging
 import org.apache.commons.io.FileUtils
 import play.api.libs.json.{Json, OFormat, Reads}
 import software.amazon.awssdk.auth.credentials.{AwsBasicCredentials, StaticCredentialsProvider}
-import software.amazon.awssdk.core.async.AsyncRequestBody
 import software.amazon.awssdk.core.checksums.RequestChecksumCalculation
 import software.amazon.awssdk.regions.Region
 import software.amazon.awssdk.services.s3.S3AsyncClient
@@ -40,9 +39,10 @@ import software.amazon.awssdk.services.s3.model.{
   CompletedMultipartUpload,
   CompletedPart,
   CreateMultipartUploadRequest,
-  PutObjectRequest,
   UploadPartRequest
 }
+import software.amazon.awssdk.transfer.s3.S3TransferManager
+import software.amazon.awssdk.transfer.s3.model.UploadDirectoryRequest
 
 import java.io.{File, RandomAccessFile}
 import java.net.URI
@@ -584,7 +584,6 @@ class UploadService @Inject()(dataSourceService: DataSourceService,
             dataSource,
             newBasePath = s"s3://$endPointHost/${dataStoreConfig.Datastore.S3Upload.bucketName}/$s3ObjectKey")
           _ <- remoteWebknossosClient.updateDataSource(s3DataSource, datasetId, allowNewPaths = true)
-          // TODO: Is uploaded dataset size the same as local dataset size?
           datasetSize <- tryo(FileUtils.sizeOfDirectoryAsBigInteger(new File(unpackToDir.toString)).longValue).toFox
           _ = this.synchronized {
             PathUtils.deleteDirectoryRecursively(unpackToDir)
@@ -671,38 +670,22 @@ class UploadService @Inject()(dataSourceService: DataSourceService,
         exploreLocalLayerService.writeLocalDatasourceProperties(dataSource, path))
     } yield path
 
+  private lazy val transferManager = S3TransferManager.builder().s3Client(s3Client).build()
+
   private def uploadDirectoryToS3(
       dataDir: Path,
       bucketName: String,
       prefix: String
   ): Fox[Unit] =
     for {
-      files <- PathUtils.listFilesRecursive(dataDir, silent = false, maxDepth = 20).toFox
-      uploadFoxes = files.map(filePath => {
-        val relPath = dataDir.relativize(filePath).toString.replace("\\", "/")
-        val s3Key = s"$prefix$relPath"
-
-        // TODO: For large files, consider using multipart upload
-        logger.info("Uploading file to S3: " + filePath)
-        val bytes = Files.readAllBytes(filePath)
-        logger.info(s"Dataset Upload: Uploading ${bytes.length} bytes to s3://$bucketName/$s3Key")
-        val startTime = System.currentTimeMillis()
-        logger.info(s"Starting upload of $filePath to S3 at $startTime")
-
-        for {
-          _ <- Fox.fromFuture {
-            s3Client
-              .putObject(
-                PutObjectRequest.builder().bucket(bucketName).key(s3Key).build(),
-                AsyncRequestBody.fromBytes(bytes)
-              )
-              .asScala
-          } ?~> s"Failed to upload file $filePath to S3"
-        } yield ()
-      })
-      // TODO: Limit number of concurrent uploads?
-      _ <- Fox.combined(uploadFoxes)
-      _ = logger.info(s"Finished uploading directory to S3 at ${System.currentTimeMillis()}")
+      _ <- Fox.successful(())
+      directoryUpload = transferManager.uploadDirectory(
+        UploadDirectoryRequest.builder().bucket(bucketName).s3Prefix(prefix).source(dataDir).build()
+      )
+      completedUpload <- Fox.fromFuture(directoryUpload.completionFuture().asScala)
+      failedTransfers = completedUpload.failedTransfers()
+      _ <- Fox.fromBool(failedTransfers.isEmpty) ?~>
+        s"Some files failed to upload to S3: $failedTransfers"
     } yield ()
 
   private def cleanUpOnFailure[T](result: Box[T],