Add support for action chaining, useful for splitting actions into ca…

…cheable pieces but still being treated as a single unit in terms of failures and logs

Protos/BuildSystem/Execution/action.proto

@@ -28,6 +28,17 @@ message LLBActionKey {
         // A merge trees based action key.
         LLBMergeTreesAction mergeTrees = 2;
     }
+
+    // A chained input is an input into the action (which may be repeated in each of the inputs for the action types)
+    // as a way to propagate information downstream for an artifact that goes through stages of processing, so that the
+    // logs for that artifact are considered a concatenation of the logs of the chain. This also groups the actions
+    // in a way where if any of the nodes in the action chain fail, then all of the downstream actions are considered
+    // as a failure, and not a dependency failure. In effect is a way to have split actions for better caching but still
+    // be considered as a single entity. There might be multiple chains sharing upstream actions, but since it's a DAG
+    // they won't be affected by the other chains, only by the upstream nodes in the chain.
+    // This field is mostly expected to be empty unless this functionality is requires, so it should be checked before
+    // being read.
+    LLBArtifact chainedInput = 3;
 }
 
 // The value for an ActionKey.

Protos/BuildSystem/Execution/action_execution.proto

@@ -27,6 +27,11 @@ message LLBActionExecutionKey {
         // A merge trees based action execution key.
         LLBMergeTreesActionExecution mergeTrees = 17;
     }
+
+    // This pairs up with the chainedInput in ActionKey. This should be used to prepopulate the logs of the action,
+    // since it contains the accumulated logs of the previous actions in the chain. This field is only present if the
+    // chainedInput field was set in the corresponding ActionKey.
+    LLBDataID chainedLogsID = 18;
 }
 
 // The value for an ActionExecutionKey.

Protos/EngineProtocol/action_execution.proto

@@ -113,6 +113,9 @@ message LLBActionExecutionRequest {
 
     // Any container for moving around unspecified data.
     repeated google.protobuf.Any additionalData = 5;
+
+    // May contain a data ID to use as base logs to provide as part of this action.
+    LLBDataID baseLogsID = 6;
 }
 
 // The response for a remote action execution request.

Sources/LLBBuildSystem/Functions/Execution/Action.swift

@@ -16,6 +16,7 @@ public extension LLBActionKey {
     static func command(
         actionSpec: LLBActionSpec,
         inputs: [LLBArtifact],
+        chainedInput: LLBArtifact? = nil,
         outputs: [LLBActionOutput],
         unconditionalOutputs: [LLBActionOutput] = [],
         mnemonic: String,
@@ -40,14 +41,23 @@ public extension LLBActionKey {
                     $0.label = label
                 }
             })
+            if let chainedInput = chainedInput {
+                $0.chainedInput = chainedInput
+            }
         }
     }
 
-    static func mergeTrees(inputs: [(artifact: LLBArtifact, path: String?)]) -> Self {
+    static func mergeTrees(
+        inputs: [(artifact: LLBArtifact, path: String?)],
+        chainedInput: LLBArtifact? = nil
+    ) -> Self {
         return LLBActionKey.with {
             $0.actionType = .mergeTrees(LLBMergeTreesAction.with {
                 $0.inputs = inputs.map { LLBMergeTreesActionInput(artifact: $0.artifact, path: $0.path) }
             })
+            if let chainedInput = chainedInput {
+                $0.chainedInput = chainedInput
+            }
         }
     }
 }
@@ -116,16 +126,32 @@ extension LLBActionKey: LLBBuildEventActionDescription {
     }
 }
 
+enum LLBChainedInputResult {
+    case none
+    case error(Error)
+    case success(LLBDataID)
+}
+
 final class ActionFunction: LLBBuildFunction<LLBActionKey, LLBActionValue> {
     override func evaluate(
         key actionKey: LLBActionKey,
         _ fi: LLBBuildFunctionInterface,
         _ ctx: Context
     ) -> LLBFuture<LLBActionValue> {
+        let chainedLogsID: LLBFuture<LLBDataID?>
+        if actionKey.hasChainedInput {
+            // Not using requestInputs since we specifically don't want the dependencyFailure processing.
+            chainedLogsID = fi.requestArtifact(actionKey.chainedInput, ctx).map { value in
+                return value.hasLogsID ? value.logsID : .none
+            }
+        } else {
+            chainedLogsID = ctx.group.next().makeSucceededFuture(.none)
+        }
+
         switch actionKey.actionType {
         case let .command(commandKey):
             ctx.buildEventDelegate?.actionScheduled(action: actionKey)
-            let resultFuture = evaluate(commandKey: commandKey, fi, ctx)
+            let resultFuture = evaluate(commandKey: commandKey, chainedLogsID: chainedLogsID, fi, ctx)
             return LLBFuture.whenAllComplete([resultFuture], on: ctx.group.next()).flatMapThrowing { results in
                 switch results[0] {
                 case .success(let value):
@@ -142,25 +168,27 @@ final class ActionFunction: LLBBuildFunction<LLBActionKey, LLBActionValue> {
                 return try results[0].get()
             }
         case let .mergeTrees(mergeTreesKey):
-            return evaluate(mergeTreesKey: mergeTreesKey, fi, ctx)
+            return evaluate(mergeTreesKey: mergeTreesKey, chainedLogsID: chainedLogsID, fi, ctx)
         case .none:
             return ctx.group.next().makeFailedFuture(LLBActionError.invalid)
         }
     }
 
     private func evaluate(
         commandKey: LLBCommandAction,
+        chainedLogsID: LLBFuture<LLBDataID?>,
         _ fi: LLBBuildFunctionInterface,
         _ ctx: Context
     ) -> LLBFuture<LLBActionValue> {
-        return fi.requestInputs(commandKey.inputs, ctx)
-            .flatMap { (inputs: [(LLBArtifact, LLBArtifactValue)]) -> LLBFuture<LLBActionValue> in
+        return chainedLogsID.and(fi.requestInputs(commandKey.inputs, ctx))
+            .flatMap { (chainedLogsID: LLBDataID?, inputs: [(LLBArtifact, LLBArtifactValue)]) -> LLBFuture<LLBActionValue> in
                 let actionExecutionKey = LLBActionExecutionKey.command(
                     actionSpec: commandKey.actionSpec,
                     inputs: inputs.map { (artifact, artifactValue) in
                         LLBActionInput(path: artifact.path, dataID: artifactValue.dataID, type: artifact.type)
                     },
                     outputs: commandKey.outputs,
+                    chainedLogsID: chainedLogsID,
                     unconditionalOutputs: commandKey.unconditionalOutputs,
                     mnemonic: commandKey.mnemonic,
                     description: commandKey.description_p,
@@ -187,13 +215,13 @@ final class ActionFunction: LLBBuildFunction<LLBActionKey, LLBActionValue> {
 
     private func evaluate(
         mergeTreesKey: LLBMergeTreesAction,
+        chainedLogsID: LLBFuture<LLBDataID?>,
         _ fi: LLBBuildFunctionInterface,
         _ ctx: Context
     ) -> LLBFuture<LLBActionValue> {
-        return fi.requestInputs(
-            mergeTreesKey.inputs.map(\.artifact),
-            ctx
-        ).flatMap { (inputs: [(artifact: LLBArtifact, artifactValue: LLBArtifactValue)]) -> LLBFuture<LLBActionExecutionValue> in
+        return chainedLogsID.and(
+            fi.requestInputs(mergeTreesKey.inputs.map(\.artifact), ctx)
+        ).flatMap { (chainedLogsID: LLBDataID?, inputs: [(artifact: LLBArtifact, artifactValue: LLBArtifactValue)]) -> LLBFuture<LLBActionExecutionValue> in
             var actionInputs = [LLBActionInput]()
             for (index, input) in mergeTreesKey.inputs.enumerated() {
                 guard inputs[index].artifact.type == .directory || !input.path.isEmpty else {
@@ -211,7 +239,10 @@ final class ActionFunction: LLBBuildFunction<LLBActionKey, LLBActionValue> {
                 )
             }
 
-            let actionExecutionKey = LLBActionExecutionKey.mergeTrees(inputs: actionInputs)
+            let actionExecutionKey = LLBActionExecutionKey.mergeTrees(
+                inputs: actionInputs,
+                chainedLogsID: chainedLogsID
+            )
 
             return fi.request(actionExecutionKey, ctx)
         }.map { actionExecutionValue in

Sources/LLBBuildSystem/Functions/Execution/ActionExecution.swift

@@ -21,6 +21,7 @@ public extension LLBActionExecutionKey {
         workingDirectory: String? = nil,
         inputs: [LLBActionInput],
         outputs: [LLBActionOutput],
+        chainedLogsID: LLBDataID? = nil,
         unconditionalOutputs: [LLBActionOutput] = [],
         mnemonic: String = "",
         description: String = "",
@@ -45,13 +46,17 @@ public extension LLBActionExecutionKey {
                 $0.description_p = description
                 $0.cacheableFailure = cacheableFailure
             })
+            if let chainedLogsID = chainedLogsID {
+                $0.chainedLogsID = chainedLogsID
+            }
         }
     }
 
     static func command(
         actionSpec: LLBActionSpec,
         inputs: [LLBActionInput],
         outputs: [LLBActionOutput],
+        chainedLogsID: LLBDataID? = nil,
         unconditionalOutputs: [LLBActionOutput] = [],
         mnemonic: String,
         description: String,
@@ -75,14 +80,20 @@ public extension LLBActionExecutionKey {
                     $0.label = label
                 }
             })
+            if let chainedLogsID = chainedLogsID {
+                $0.chainedLogsID = chainedLogsID
+            }
         }
     }
 
-    static func mergeTrees(inputs: [LLBActionInput]) -> Self {
+    static func mergeTrees(inputs: [LLBActionInput], chainedLogsID: LLBDataID? = nil) -> Self {
         return LLBActionExecutionKey.with {
             $0.actionExecutionType = .mergeTrees(LLBMergeTreesActionExecution.with {
                 $0.inputs = inputs
             })
+            if let chainedLogsID = chainedLogsID {
+                $0.chainedLogsID = chainedLogsID
+            }
         }
     }
 }
@@ -160,7 +171,17 @@ final class ActionExecutionFunction: LLBBuildFunction<LLBActionExecutionKey, LLB
         self.dynamicActionExecutorDelegate = dynamicActionExecutorDelegate
     }
 
-    override func evaluate(key actionExecutionKey: LLBActionExecutionKey, _ fi: LLBBuildFunctionInterface, _ ctx: Context) -> LLBFuture<LLBActionExecutionValue> {
+    override func evaluate(
+        key actionExecutionKey: LLBActionExecutionKey,
+        _ fi: LLBBuildFunctionInterface,
+        _ ctx: Context
+    ) -> LLBFuture<LLBActionExecutionValue> {
+        let chainedLogsID: LLBDataID?
+        if actionExecutionKey.hasChainedLogsID {
+            chainedLogsID = actionExecutionKey.chainedLogsID
+        } else {
+            chainedLogsID = nil
+        }
 
         switch actionExecutionKey.actionExecutionType {
         case let .command(commandKey):
@@ -170,22 +191,29 @@ final class ActionExecutionFunction: LLBBuildFunction<LLBActionExecutionKey, LLB
                 description: actionExecutionKey.description,
                 owner: actionExecutionKey.owner
             )
-            return evaluateCommand(commandKey: commandKey, requestExtras: requestExtras, fi, ctx).map {
+            return evaluateCommand(
+                commandKey: commandKey,
+                chainedLogsID: chainedLogsID,
+                requestExtras: requestExtras,
+                fi,
+                ctx
+            ).map {
                 ctx.buildEventDelegate?.actionExecutionCompleted(action: actionExecutionKey)
                 return $0
             }.flatMapErrorThrowing { error in
                 ctx.buildEventDelegate?.actionExecutionCompleted(action: actionExecutionKey)
                 throw error
             }
         case let .mergeTrees(mergeTreesKey):
-            return evaluateMergeTrees(mergeTreesKey: mergeTreesKey, fi, ctx)
+            return evaluateMergeTrees(mergeTreesKey: mergeTreesKey, chainedLogsID: chainedLogsID, fi, ctx)
         case .none:
             return ctx.group.next().makeFailedFuture(LLBActionExecutionError.invalid)
         }
     }
 
     private func evaluateCommand(
         commandKey: LLBCommandActionExecution,
+        chainedLogsID: LLBDataID?,
         requestExtras: LLBActionExecutionRequestExtras,
         _ fi: LLBBuildFunctionInterface,
         _ ctx: Context
@@ -203,7 +231,8 @@ final class ActionExecutionFunction: LLBBuildFunction<LLBActionExecutionKey, LLB
             inputs: commandKey.inputs,
             outputs: commandKey.outputs,
             unconditionalOutputs: commandKey.unconditionalOutputs,
-            additionalData: additionalRequestData
+            additionalData: additionalRequestData,
+            baseLogsID: chainedLogsID
         )
 
         let resultFuture: LLBFuture<LLBActionExecutionResponse>
@@ -239,7 +268,12 @@ final class ActionExecutionFunction: LLBBuildFunction<LLBActionExecutionKey, LLB
         }
     }
 
-    private func evaluateMergeTrees(mergeTreesKey: LLBMergeTreesActionExecution, _ fi: LLBBuildFunctionInterface, _ ctx: Context) -> LLBFuture<LLBActionExecutionValue> {
+    private func evaluateMergeTrees(
+        mergeTreesKey: LLBMergeTreesActionExecution,
+        chainedLogsID: LLBDataID?,
+        _ fi: LLBBuildFunctionInterface,
+        _ ctx: Context
+    ) -> LLBFuture<LLBActionExecutionValue> {
         let inputs = mergeTreesKey.inputs
         // Skip merging if there's a single tree as input, with no path to prepend.
         if inputs.count == 1, inputs[0].type == .directory, inputs[0].path.isEmpty {
@@ -256,7 +290,9 @@ final class ActionExecutionFunction: LLBBuildFunction<LLBActionExecutionKey, LLB
 
         // Skip merging if there is a single prepended tree.
         if prependedTrees.count == 1 {
-            return prependedTrees[0].map { LLBActionExecutionValue(outputs: [$0.id], stdoutID: nil, stderrID: nil) }
+            return prependedTrees[0].map {
+                LLBActionExecutionValue(outputs: [$0.id], stdoutID: chainedLogsID, stderrID: chainedLogsID)
+            }
         }
 
         return LLBFuture.whenAllSucceed(prependedTrees, on: ctx.group.next()).flatMap { trees in