Merge pull request #40 from Automattic/isolate/s3

Reduce SotoS3 dependency

Author: jkmassel

Package.resolved

@@ -17,8 +17,8 @@ let package = Package(
         .package(url: "https://github.com/apple/swift-log.git", from: "1.0.0"),
         .package(url: "https://github.com/jkmassel/kcpassword-swift.git", from: "1.0.0"),
         .package(url: "https://github.com/swiftpackages/DotEnv.git", from: "3.0.0"),
-        .package(url: "https://github.com/apple/swift-tools-support-core", from: "0.2.5")
-
+        .package(url: "https://github.com/apple/swift-tools-support-core", from: "0.2.5"),
+        .package(url: "https://github.com/Alamofire/Alamofire.git", .upToNextMajor(from: "5.6.1"))
     ],
     targets: [
         // Targets are the basic building blocks of a package. A target can define a module or a test suite.
@@ -27,20 +27,20 @@ let package = Package(
             name: "hostmgr",
             dependencies: [
                 .product(name: "ArgumentParser", package: "swift-argument-parser"),
-                .product(name: "SotoS3", package: "soto"),
                 .product(name: "prlctl", package: "prlctl"),
                 .product(name: "Tqdm", package: "swift-tqdm"),
                 .product(name: "Logging", package: "swift-log"),
                 .product(name: "kcpassword", package: "kcpassword-swift"),
-                .target(name: "libhostmgr")
+                .target(name: "libhostmgr"),
             ]
         ),
         .target(
             name: "libhostmgr",
             dependencies: [
                 .product(name: "ArgumentParser", package: "swift-argument-parser"),
                 .product(name: "SotoS3", package: "soto"),
-                .product(name: "TSCBasic", package: "swift-tools-support-core")
+                .product(name: "TSCBasic", package: "swift-tools-support-core"),
+                .product(name: "Alamofire", package: "Alamofire"),
             ]
         ),
         .testTarget(

Package.swift

@@ -1,7 +1,7 @@
 import Foundation
 import ArgumentParser
 
-struct BenchmarkCommand: ParsableCommand {
+struct BenchmarkCommand: AsyncParsableCommand {
     static let configuration = CommandConfiguration(
         commandName: "benchmark",
         abstract: "System tests",

Sources/hostmgr/BenchmarkCommand.swift

@@ -1,5 +1,4 @@
 import ArgumentParser
-import SotoS3
 import Logging
 import libhostmgr
 
@@ -66,12 +65,10 @@ struct InitCommand: ParsableCommand {
             currentValue: Configuration.shared.vmImagesBucket
         )
 
-        let vmImagesRegion = prompt(
+        Configuration.shared.vmImagesRegion = prompt(
             "Which AWS region contains the \(Configuration.shared.vmImagesBucket) bucket?",
-            currentValue: Configuration.shared.vmImagesRegion.rawValue
-        ) { Region(awsRegionName: $0) != nil }
-
-        Configuration.shared.vmImagesRegion = Region(awsRegionName: vmImagesRegion)!
+            currentValue: Configuration.shared.vmImagesRegion
+        )
 
         print("== Authorized Keys Sync ==\n")
 
@@ -80,12 +77,10 @@ struct InitCommand: ParsableCommand {
             currentValue: Configuration.shared.authorizedKeysBucket
         )
 
-        let authorizedKeysRegion = prompt(
+        Configuration.shared.authorizedKeysRegion = prompt(
             "Which AWS region contains the \(Configuration.shared.authorizedKeysBucket) bucket?",
-            currentValue: Configuration.shared.authorizedKeysRegion.rawValue
-        ) { Region(awsRegionName: $0) != nil }
-
-        Configuration.shared.authorizedKeysRegion = Region(awsRegionName: authorizedKeysRegion)!
+            currentValue: Configuration.shared.authorizedKeysRegion
+        )
 
         Configuration.shared.gitMirrorBucket = prompt(
             "Which S3 bucket would you like to use as the data source for the git mirror server?",

Sources/hostmgr/HostMgrCommand.swift

@@ -2,7 +2,7 @@ import Foundation
 import ArgumentParser
 import libhostmgr
 
-struct SyncCommand: ParsableCommand {
+struct SyncCommand: AsyncParsableCommand {
     static let configuration = CommandConfiguration(
         commandName: "sync",
         abstract: "Sync remote data with this host",
@@ -18,7 +18,7 @@ struct SyncCommand: ParsableCommand {
     @OptionGroup
     var options: SharedSyncOptions
 
-    func run() throws {
+    mutating func run() async throws {
 
         if list {
             Configuration.SchedulableSyncCommand.allCases.forEach { print($0) }
@@ -29,20 +29,20 @@ struct SyncCommand: ParsableCommand {
         // to make it so it doesn't need to be installed separately
         try GenerateGitMirrorManifestTask().run()
 
-        try Configuration.shared.syncTasks.forEach { command in
-            options.force ? print("Force-running \(command.rawValue)") : print("Running \(command.rawValue)")
-            try perform(task: command, immediately: options.force)
+        for task in Configuration.shared.syncTasks {
+            options.force ? print("Force-running \(task.rawValue)") : print("Running \(task.rawValue)")
+            try await perform(task: task)
         }
     }
 
-    private func perform(task: Configuration.SchedulableSyncCommand, immediately: Bool) throws {
+    private func perform(task: Configuration.SchedulableSyncCommand) async throws {
         switch task {
         case .authorizedKeys:
             let command = SyncAuthorizedKeysCommand(options: self._options)
-            try command.run()
+            try await command.run()
         case .vmImages:
             let command = SyncVMImagesCommand(options: self._options)
-            try command.run()
+            try await command.run()
         }
     }
 }

Sources/hostmgr/SyncCommand.swift

@@ -2,7 +2,7 @@ import Foundation
 import ArgumentParser
 import prlctl
 
-struct VMCommand: ParsableCommand {
+struct VMCommand: AsyncParsableCommand {
 
     static let configuration = CommandConfiguration(
         commandName: "vm",

Sources/hostmgr/VMCommand.swift

@@ -6,41 +6,53 @@ import libhostmgr
 
 private let startDate = Date()
 
-struct NetworkBenchmark: ParsableCommand {
+struct NetworkBenchmark: AsyncParsableCommand {
 
     static let configuration = CommandConfiguration(
         commandName: "network",
         abstract: "Test Network Speed"
     )
 
-    func run() throws {
-        guard let file = try VMRemoteImageManager().list().sorted(by: { $0.size < $1.size }).first else {
+    private static let limiter = Limiter(policy: .throttle, operationsPerSecond: 1)
+
+    func run() async throws {
+        let remoteImages = try await RemoteVMRepository().listImages(sortedBy: .size)
+
+        guard let file = remoteImages.last else {
             throw CleanExit.message("Unable to find a remote image to use as a network benchmark")
         }
 
-        try S3Manager().streamingDownloadFile(
-            region: Configuration.shared.vmImagesRegion,
+        let manager = S3Manager(
             bucket: Configuration.shared.vmImagesBucket,
-            key: file.imageKey,
-            destination: URL(fileURLWithPath: "/dev/null"),
-            progressCallback: self.showProgress
+            region: Configuration.shared.vmImagesRegion
+        )
+
+        try await manager.download(
+            object: file.imageObject,
+            to: FileManager.default.temporaryFilePath(),
+            progressCallback: self.updateProgress
         )
     }
 
-    private func showProgress(availableBytes: Int, downloadedBytes: Int, totalBytes: Int64) {
+    private func imageSizeSort(_ lhs: RemoteVMImage, _ rhs: RemoteVMImage) -> Bool {
+        lhs.imageObject.size < rhs.imageObject.size
+    }
 
-        // Sample only one in 100 entries
-        guard Int.random(in: 0...1000) == 0 else {
-            return
-        }
+    private func updateProgress(_ progress: FileTransferProgress) {
+        Self.limiter.perform {
+            let downloadedSize = ByteCountFormatter.string(fromByteCount: Int64(progress.current), countStyle: .file)
+            let totalSize = ByteCountFormatter.string(fromByteCount: Int64(progress.total), countStyle: .file)
 
-        let downloadedSize = ByteCountFormatter.string(fromByteCount: Int64(downloadedBytes), countStyle: .file)
-        let totalSize = ByteCountFormatter.string(fromByteCount: totalBytes, countStyle: .file)
+            let secondsElapsed = Date().timeIntervalSince(startDate)
+            let perSecond = Double(progress.current) / Double(secondsElapsed)
 
-        let secondsElapsed = Date().timeIntervalSince(startDate)
-        let perSecond = Double(downloadedBytes) / Double(secondsElapsed)
-        let rate = ByteCountFormatter.string(fromByteCount: Int64(perSecond), countStyle: .file)
+            // Don't continue unless the rate can be represented by `Int64`
+            guard perSecond.isNormal else {
+                return
+            }
 
-        logger.info("Downloaded \(downloadedSize) of \(totalSize) [Rate: \(rate) per second]")
+            let rate = ByteCountFormatter.string(fromByteCount: Int64(perSecond), countStyle: .file)
+            logger.info("Downloaded \(downloadedSize) of \(totalSize) [Rate: \(rate) per second]")
+        }
     }
 }

Sources/hostmgr/commands/benchmark/NetworkBenchmark.swift

@@ -1,6 +1,5 @@
 import Foundation
 import ArgumentParser
-import SotoS3
 import kcpassword
 
 struct SetAutomaticLoginPasswordCommand: ParsableCommand {

Sources/hostmgr/commands/set/SetAutomaticLoginPasswordCommand.swift

@@ -1,12 +1,11 @@
 import Foundation
 import ArgumentParser
-import SotoS3
 import libhostmgr
 
-struct SyncAuthorizedKeysCommand: ParsableCommand, FollowsCommandPolicies {
+struct SyncAuthorizedKeysCommand: AsyncParsableCommand, FollowsCommandPolicies {
 
     static let configuration = CommandConfiguration(
-        commandName: "authorized_keys",
+        commandName: Configuration.SchedulableSyncCommand.authorizedKeys.rawValue,
         abstract: "Set this machine's authorized_keys file"
     )
 
@@ -20,7 +19,7 @@ struct SyncAuthorizedKeysCommand: ParsableCommand, FollowsCommandPolicies {
         name: .shortAndLong,
         help: "The S3 region for the bucket"
     )
-    var region: Region = Configuration.shared.authorizedKeysRegion
+    var region: String = Configuration.shared.authorizedKeysRegion
 
     @Option(
         name: .shortAndLong,
@@ -44,30 +43,26 @@ struct SyncAuthorizedKeysCommand: ParsableCommand, FollowsCommandPolicies {
         .scheduled(every: 3600)
     ]
 
-    func run() throws {
+    func run() async throws {
         try to(evaluateCommandPolicies(), unless: options.force)
+        logger.debug("Job schedule allows for running")
 
-        logger.debug("Downloading file from s3://\(bucket)/\(key) in \(region) to \(destination)")
-        logger.trace("Job schedule allows for running")
+        logger.info("Downloading file from s3://\(bucket)/\(key) in \(region) to \(destination)")
 
-        guard let bytes = try S3Manager().getFileBytes(region: region, bucket: bucket, key: key) else {
-            print("Unable to sync authorized_keys file – exiting")
-            SyncAuthorizedKeysCommand.exit()
-        }
-
-        logger.trace("Downloaded \(bytes.count) bytes from S3")
+        let s3Manager = S3Manager(bucket: self.bucket, region: self.region)
 
-        /// Create the parent directory if needed
-        let parent = URL(fileURLWithPath: Configuration.shared.localAuthorizedKeys).deletingLastPathComponent()
-        try FileManager.default.createDirectoryTree(atUrl: parent)
+        guard let object = try await s3Manager.lookupObject(atPath: key) else {
+            logger.error("Unable to locate authorized_keys file – exiting")
+            throw ExitCode(rawValue: 1)
+        }
 
-        /// Overwrite the existing file
-        try bytes.write(to: URL(fileURLWithPath: destination))
+        let url = URL(fileURLWithPath: self.destination)
+        try await s3Manager.download(object: object, to: url, progressCallback: nil)
 
         /// Fix the permissions on the file, if needed
         try FileManager.default.setAttributes([
             .posixPermissions: 0o600
-        ], ofItemAtPath: destination)
+        ], ofItemAtPath: self.destination)
 
         try recordLastRun()
     }

Sources/hostmgr/commands/sync/SyncAuthorizedKeysCommand.swift

@@ -1,14 +1,13 @@
 import Foundation
 import ArgumentParser
-import SotoS3
 import prlctl
 import libhostmgr
 
-struct SyncVMImagesCommand: ParsableCommand, FollowsCommandPolicies {
+struct SyncVMImagesCommand: AsyncParsableCommand, FollowsCommandPolicies {
 
     static let configuration = CommandConfiguration(
-        commandName: "vm_images",
-        abstract: "Sync this machine's VM images with those avaiable remotely"
+        commandName: Configuration.SchedulableSyncCommand.vmImages.rawValue,
+        abstract: "Sync this machine's VM images with those available remotely"
     )
 
     @OptionGroup
@@ -22,19 +21,19 @@ struct SyncVMImagesCommand: ParsableCommand, FollowsCommandPolicies {
         .scheduled(every: 3600)
     ]
 
-    func run() throws {
+    func run() async throws {
         try to(evaluateCommandPolicies(), unless: options.force)
 
         /// The manifest defines which images should be distributed to VM hosts
-        let manifest = try VMRemoteImageManager().getManifest()
-        logger.debug("Downloaded manifest:\n\(manifest)")
+        let manifest = try await RemoteVMRepository().getManifest()
+        logger.info("Downloaded manifest:\n\(manifest)")
 
         /// Candidate images are any that the manifest says *could* be installed on this VM host
-        let candidateImages = try VMRemoteImageManager().list().filter { manifest.contains($0.basename) }
-        logger.debug("Available remote images:\n\(candidateImages)")
+        let candidateImages = try await RemoteVMRepository().listImages().filter { manifest.contains($0.basename) }
+        logger.info("Available remote images:\n\(candidateImages)")
 
         let localImages = try VMLocalImageManager().list()
-        logger.debug("Local Images:\(localImages)")
+        logger.info("Local Images:\(localImages)")
 
         let imagesToDownload = candidateImages.filter { !localImages.contains($0.basename) }
         let imagesToDelete = localImages
@@ -46,27 +45,25 @@ struct SyncVMImagesCommand: ParsableCommand, FollowsCommandPolicies {
         logger.info("Deleting local images:\(imagesToDelete)")
         try VMLocalImageManager().delete(images: imagesToDelete)
 
-        try imagesToDownload.forEach {
-            try download(image: $0)
+        for image in imagesToDownload {
+            try await download(image: image)
         }
 
         try recordLastRun()
     }
 
-    private func download(image: VMRemoteImageManager.RemoteImage) throws {
+    private func download(image: RemoteVMImage) async throws {
         let storageDirectory = Configuration.shared.vmStorageDirectory
         let destination = storageDirectory.appendingPathComponent(image.fileName)
 
-        logger.info("Downloading the VM – this will take a few minutes")
-        logger.trace("Downloading \(image.basename) to \(destination)")
+        logger.info("Downloading \(image.basename) to \(destination) – this will take a few minutes")
 
         let limiter = Limiter(policy: .throttle, operationsPerSecond: 1)
 
-        try VMRemoteImageManager().download(image: image, to: destination) { _, downloaded, total in
+        try await RemoteVMRepository().download(image: image, to: destination) { progress in
             limiter.perform {
                 try? recordHeartbeat()
-                let percent = String(format: "%.2f", Double(downloaded) / Double(total) * 100)
-                logger.trace("\(percent)% complete")
+                logger.info("\(progress.decimalPercent)% complete")
             }
         }
 
@@ -87,6 +84,5 @@ struct SyncVMImagesCommand: ParsableCommand, FollowsCommandPolicies {
         logger.info("Imported Complete")
         logger.info("\tName:\t\(vmToImport.name)")
         logger.info("\tUUID:\t\(vmToImport.uuid)")
-
     }
 }