Merge pull request #20 from Automattic/fix/vm-image-sync

Fix sync

Author: jkmassel

.buildkite/pipeline.yml

@@ -4,7 +4,7 @@ common_params:
   - &bash_cache automattic/bash-cache#v1.3.2: ~
   # Common environment values to use with the `env` key.
   env: &common_env
-    IMAGE_ID: xcode-12.5.1
+    IMAGE_ID: xcode-13.4.1
 
 # This is the default pipeline – it will build and test the app
 steps:

Sources/hostmgr/SyncCommand.swift

@@ -15,19 +15,11 @@ struct SyncCommand: ParsableCommand {
     @Flag(help: "List available sync tasks")
     var list: Bool = false
 
-    @Flag(help: "Force the job to run immediately, ignoring the schedule")
-    var force: Bool = false
-
-    @Argument
-    var task: Configuration.SchedulableSyncCommand?
+    @OptionGroup
+    var options: SharedSyncOptions
 
     func run() throws {
 
-        if let task = task {
-            try perform(task: task, immediately: force)
-            return
-        }
-
         if list {
             Configuration.SchedulableSyncCommand.allCases.forEach { print($0) }
             return
@@ -38,15 +30,19 @@ struct SyncCommand: ParsableCommand {
         try GenerateGitMirrorManifestTask().run()
 
         try Configuration.shared.syncTasks.forEach { command in
-            force ? print("Force-running \(command.rawValue)") : print("Running \(command.rawValue)")
-            try perform(task: command, immediately: self.force)
+            options.force ? print("Force-running \(command.rawValue)") : print("Running \(command.rawValue)")
+            try perform(task: command, immediately: options.force)
         }
     }
 
     private func perform(task: Configuration.SchedulableSyncCommand, immediately: Bool) throws {
         switch task {
-        case .authorizedKeys: try SyncAuthorizedKeysTask().run(force: immediately)
-        case .vmImages: try SyncVMImagesTask().run(force: immediately)
+        case .authorizedKeys:
+            let command = SyncAuthorizedKeysCommand(options: self._options)
+            try command.run()
+        case .vmImages:
+            let command = SyncVMImagesCommand(options: self._options)
+            try command.run()
         }
     }
 }

Sources/hostmgr/commands/sync/SharedSyncOptions.swift

@@ -0,0 +1,6 @@
+import ArgumentParser
+
+struct SharedSyncOptions: ParsableArguments {
+    @Flag(help: "Force all jobs to run immediately, ignoring the schedule")
+    var force: Bool = false
+}

Sources/hostmgr/commands/sync/SyncAuthorizedKeysCommand.swift

@@ -3,7 +3,7 @@ import ArgumentParser
 import SotoS3
 import libhostmgr
 
-struct SyncAuthorizedKeysCommand: ParsableCommand {
+struct SyncAuthorizedKeysCommand: ParsableCommand, FollowsCommandPolicies {
 
     static let configuration = CommandConfiguration(
         commandName: "authorized_keys",
@@ -34,40 +34,20 @@ struct SyncAuthorizedKeysCommand: ParsableCommand {
     )
     var destination: String = Configuration.shared.localAuthorizedKeys
 
-    func run() throws {
-        try SyncAuthorizedKeysTask(bucket: bucket, region: region, key: key, destination: destination).run()
-    }
-}
-
-struct SyncAuthorizedKeysTask {
-
-    private let bucket: String
-    private let region: Region
-    private let key: String
-    private let destination: String
-
-    init(
-        bucket: String = Configuration.shared.authorizedKeysBucket,
-        region: Region = Configuration.shared.authorizedKeysRegion,
-        key: String = "authorized_keys",
-        destination: String = Configuration.shared.localAuthorizedKeys
-    ) {
-        self.bucket = bucket
-        self.region = region
-        self.key = key
-        self.destination = destination
-    }
+    @OptionGroup
+    var options: SharedSyncOptions
 
-    func run(force: Bool = false) throws {
-        let state = State.get()
+    static let commandIdentifier: String = "authorized-key-sync"
 
-        logger.debug("Downloading file from s3://\(bucket)/\(key) in \(region) to \(destination)")
+    /// A set of command policies that control the circumstances under which this command can be run
+    static let commandPolicies: [CommandPolicy] = [
+        .scheduled(every: 3600)
+    ]
 
-        guard state.shouldRun && force else {
-            print("This job is not scheduled to run until \(state.nextRunTime)")
-            return
-        }
+    func run() throws {
+        try to(evaluateCommandPolicies(), unless: options.force)
 
+        logger.debug("Downloading file from s3://\(bucket)/\(key) in \(region) to \(destination)")
         logger.trace("Job schedule allows for running")
 
         guard let bytes = try S3Manager().getFileBytes(region: region, bucket: bucket, key: key) else {
@@ -89,29 +69,6 @@ struct SyncAuthorizedKeysTask {
             .posixPermissions: 0o600
         ], ofItemAtPath: destination)
 
-        try State.set(state: State(lastRunAt: Date()))
-    }
-
-    struct State: Codable {
-        private static let key = "authorized-key-sync-state"
-        var lastRunAt: Date = Date.distantPast
-
-        var shouldRun: Bool {
-            let runInterval = TimeInterval(Configuration.shared.authorizedKeysSyncInterval)
-            return self.lastRunAt < Date().addingTimeInterval(runInterval * -1)
-        }
-
-        var nextRunTime: Date {
-            let runInterval = TimeInterval(Configuration.shared.authorizedKeysSyncInterval)
-            return self.lastRunAt.addingTimeInterval(runInterval)
-        }
-
-        static func get() -> State {
-            (try? StateManager.load(key: key)) ?? State()
-        }
-
-        static func set(state: State) throws {
-            try StateManager.store(key: key, value: state)
-        }
+        try recordLastRun()
     }
 }

Sources/hostmgr/commands/sync/SyncVMImagesCommand.swift

@@ -4,32 +4,26 @@ import SotoS3
 import prlctl
 import libhostmgr
 
-struct SyncVMImagesCommand: ParsableCommand {
+struct SyncVMImagesCommand: ParsableCommand, FollowsCommandPolicies {
 
     static let configuration = CommandConfiguration(
         commandName: "vm_images",
         abstract: "Sync this machine's VM images with those avaiable remotely"
     )
 
-    func run() throws {
-        try SyncVMImagesTask().run()
-    }
-}
+    @OptionGroup
+    var options: SharedSyncOptions
 
-struct SyncVMImagesTask {
+    static let commandIdentifier: String = "sync-vm-images"
 
-    let storageDirectory = Configuration.shared.vmStorageDirectory
+    /// A set of command policies that control the circumstances under which this command can be run
+    static let commandPolicies: [CommandPolicy] = [
+        .serialExecution,
+        .scheduled(every: 3600)
+    ]
 
-    func run(force: Bool = false, state: State = State.get()) throws {
-        guard !state.isRunning else {
-            print("Already syncing VMs, so we won't try to run again")
-            return
-        }
-
-        guard state.shouldRun && force else {
-            print("This job is not scheduled to run until \(state.nextRunTime)")
-            return
-        }
+    func run() throws {
+        try to(evaluateCommandPolicies(), unless: options.force)
 
         /// The manifest defines which images should be distributed to VM hosts
         let manifest = try VMRemoteImageManager().getManifest()
@@ -53,31 +47,27 @@ struct SyncVMImagesTask {
         try VMLocalImageManager().delete(images: imagesToDelete)
 
         try imagesToDownload.forEach {
-            try download(image: $0, state: state)
+            try download(image: $0)
         }
 
-        try State.set(state: State(lastRunAt: Date()))
+        try recordLastRun()
     }
 
-    private func download(image: VMRemoteImageManager.RemoteImage, state: State = State.get()) throws {
+    private func download(image: VMRemoteImageManager.RemoteImage) 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)")
 
+        let limiter = Limiter(policy: .throttle, operationsPerSecond: 1)
+
         try VMRemoteImageManager().download(image: image, to: destination) { _, downloaded, total in
-            /// Only update the heartbeat every 5 seconds to avoid thrashing the disk
-            guard abs(state.heartBeat.timeIntervalSinceNow) > 5 else {
-                return
+            limiter.perform {
+                try? recordHeartbeat()
+                let percent = String(format: "%.2f", Double(downloaded) / Double(total) * 100)
+                logger.trace("\(percent)% complete")
             }
-
-            try? State.set(state: State(
-                lastRunAt: state.lastRunAt,
-                heartBeat: Date()
-            ))
-
-            let percent = String(format: "%.2f", Double(downloaded) / Double(total) * 100)
-            logger.trace("\(percent)% complete")
         }
 
         logger.info("Download Complete")
@@ -99,34 +89,4 @@ struct SyncVMImagesTask {
         logger.info("\tUUID:\t\(vmToImport.uuid)")
 
     }
-
-    struct State: Codable {
-        private static let key = "sync-vm-images-state"
-        var lastRunAt: Date = Date.distantPast
-
-        var shouldRun: Bool {
-            let runInterval = TimeInterval(Configuration.shared.authorizedKeysSyncInterval)
-            return self.lastRunAt < Date().addingTimeInterval(runInterval * -1)
-        }
-
-        var nextRunTime: Date {
-            let runInterval = TimeInterval(Configuration.shared.authorizedKeysSyncInterval)
-            return self.lastRunAt.addingTimeInterval(runInterval)
-        }
-
-        var heartBeat: Date = Date.distantPast
-
-        // if we haven't hear from a job in 60 seconds, assume it's failed and we should try again
-        var isRunning: Bool {
-            abs(heartBeat.timeIntervalSinceNow) < 60
-        }
-
-        static func get() -> State {
-            (try? StateManager.load(key: key)) ?? State()
-        }
-
-        static func set(state: State) throws {
-            try StateManager.store(key: key, value: state)
-        }
-    }
 }

Sources/hostmgr/helpers/Foundation.swift

@@ -1,32 +1,6 @@
 import Foundation
 import ArgumentParser
-
-extension FileManager {
-
-    func createTemporaryFile(containing string: String = "") throws -> URL {
-        let file = temporaryDirectory.appendingPathComponent(UUID().uuidString)
-        try string.write(to: file, atomically: false, encoding: .utf8)
-        return file
-    }
-
-    func directoryExists(atUrl url: URL) -> Bool {
-        var isDirectory: ObjCBool = false
-        let exists = self.fileExists(atPath: url.path, isDirectory: &isDirectory)
-        return exists && isDirectory.boolValue
-    }
-
-    func createDirectoryTree(atUrl url: URL) throws {
-        try FileManager.default.createDirectory(at: url, withIntermediateDirectories: true)
-    }
-
-    func subpaths(at url: URL) -> [String] {
-        self.subpaths(atPath: url.path) ?? []
-    }
-
-    func displayName(at url: URL) -> String {
-        displayName(atPath: url.path)
-    }
-}
+import libhostmgr
 
 extension URL: ExpressibleByArgument {
     public init?(argument: String) {

Sources/hostmgr/helpers/S3.swift

@@ -91,7 +91,7 @@ struct S3Manager {
         let estimatedDownloadSpeedInMBPS: Int64 = 10
         let minimalTimeout: Int64 = 60
         let timeout = max(totalMB / estimatedDownloadSpeedInMBPS, minimalTimeout)
-        logger.info("Download timeout: \(timeout / 60) minutes")
+        logger.info("Estimated Download Time: \(timeout / 60) minutes")
 
         let s3Client = try getS3Client(from: client, for: bucket, in: region).with(timeout: .seconds(timeout))
         let objectRequest = S3.GetObjectRequest(bucket: bucket, key: key)