Move reporting of result out of do_work

Summary:
Now with us reporting the result at the outermost place, we can also
report the actual CancellationReason.

Reviewed By: JakobDegen

Differential Revision: D65371045

fbshipit-source-id: 4d77c9a41b63b706b122e349c98d0599e98b6697

dice/dice/src/impls/task/dice.rs

@@ -284,7 +284,7 @@ impl DiceTaskInternal {
         }
     }
 
-    pub(super) fn set_value(
+    pub(crate) fn set_value(
         &self,
         value: DiceComputedValue,
     ) -> CancellableResult<DiceComputedValue> {
@@ -332,7 +332,7 @@ impl DiceTaskInternal {
 
     /// report the task as terminated. This should only be called once. No effect if called affect
     /// task is already ready
-    pub(super) fn report_terminated(&self, reason: CancellationReason) {
+    pub(crate) fn report_terminated(&self, reason: CancellationReason) {
         match self.state.sync() {
             TaskState::Continue => {}
             TaskState::Finished => {

dice/dice/src/impls/task/handle.rs

@@ -68,6 +68,10 @@ impl<'a> DiceTaskHandle<'a> {
     pub(crate) fn finished(&mut self, value: DiceComputedValue) {
         self.result = Some(Ok(value));
     }
+
+    pub(crate) fn cancelled(&mut self, reason: CancellationReason) {
+        self.result = Some(Err(reason));
+    }
 }
 
 unsafe impl<'a> Send for DiceTaskHandle<'a> {}

dice/dice/src/impls/worker.rs

@@ -115,12 +115,23 @@ impl DiceTaskWorker {
                     None => Either::Right(state.no_previous_task(handle).await),
                 };
 
-                match previous_result {
-                    Either::Left(_) => {
+                let result = match previous_result {
+                    Either::Left(previous_result) => {
                         // previous result actually finished
+                        previous_result
                     }
-                    Either::Right(state) => {
-                        let _ignore = worker.do_work(handle, state_handle, state).await;
+                    Either::Right(state) => worker.do_work(handle, state_handle, state).await,
+                };
+
+                match result {
+                    Ok(DiceWorkerStateFinishedAndCached {
+                        _prevent_cancellation,
+                        value,
+                    }) => {
+                        handle.finished(value);
+                    }
+                    Err(reason) => {
+                        handle.cancelled(reason);
                     }
                 }
 
@@ -208,7 +219,7 @@ impl DiceTaskWorker {
                             )
                             .await;
 
-                        return response.map(|r| task_state.cached(handle, r, activation_info));
+                        return response.map(|r| task_state.cached(r, activation_info));
                     }
                     CheckDependenciesResult::NoDeps => {
                         // TODO(cjhopman): Why do we treat nodeps as deps not matching? There seems to be some
@@ -265,7 +276,7 @@ impl DiceTaskWorker {
             }
         };
 
-        res.map(|res| state.cached(handle, res, activation_info))
+        res.map(|res| state.cached(res, activation_info))
     }
 
     async fn compute(

dice/dice/src/impls/worker/state.rs

@@ -56,14 +56,12 @@ impl<'a> DiceWorkerStateAwaitingPrevious<'a> {
 
     pub(crate) fn previously_finished(
         self,
-        internals: &mut DiceTaskHandle,
         value: DiceComputedValue,
     ) -> CancellableResult<DiceWorkerStateFinishedAndCached> {
         debug!(msg = "previously cancelled task actually finished");
 
         let guard = self.prevent_cancellation.try_disable_cancellation();
-
-        finish_with_cached_value(internals, value, guard)
+        finish_with_cached_value(value, guard)
     }
 
     pub(crate) async fn previously_cancelled(
@@ -115,7 +113,7 @@ impl<'a> DiceWorkerStateAwaitingPrevious<'a> {
             .expect("Terminated task must have finished value")
         {
             Ok(res) => {
-                return Either::Left(self.previously_finished(internals, res));
+                return Either::Left(self.previously_finished(res));
             }
             Err(_cancelled) => {
                 // actually was cancelled, so just continue re-evaluating
@@ -127,19 +125,15 @@ impl<'a> DiceWorkerStateAwaitingPrevious<'a> {
 }
 
 fn finish_with_cached_value(
-    internals: &mut DiceTaskHandle<'_>,
     value: DiceComputedValue,
     disable_cancellation: Option<DisableCancellationGuard>,
 ) -> CancellableResult<DiceWorkerStateFinishedAndCached> {
     match disable_cancellation {
         None => Err(CancellationReason::Cached),
-        Some(g) => {
-            internals.finished(value);
-
-            Ok(DiceWorkerStateFinishedAndCached {
-                _prevent_cancellation: g,
-            })
-        }
+        Some(g) => Ok(DiceWorkerStateFinishedAndCached {
+            value,
+            _prevent_cancellation: g,
+        }),
     }
 }
 
@@ -198,7 +192,7 @@ impl DiceWorkerStateLookupNode {
         debug!(msg = "found existing entry with matching version in cache. reusing result.");
 
         let guard = internals.cancellation_ctx().try_disable_cancellation();
-        finish_with_cached_value(internals, value, guard)
+        finish_with_cached_value(value, guard)
     }
 }
 
@@ -287,7 +281,6 @@ pub(crate) struct DiceWorkerStateFinished {
 impl DiceWorkerStateFinished {
     pub(crate) fn cached(
         self,
-        internals: &mut DiceTaskHandle,
         value: DiceComputedValue,
         activation_info: Option<ActivationInfo>,
     ) -> DiceWorkerStateFinishedAndCached {
@@ -300,9 +293,9 @@ impl DiceWorkerStateFinished {
                 activation_info.activation_data,
             )
         }
-        internals.finished(value);
 
         DiceWorkerStateFinishedAndCached {
+            value,
             _prevent_cancellation: self._prevent_cancellation,
         }
     }
@@ -342,5 +335,6 @@ impl ActivationInfo {
 /// When the spawned dice worker is done computing and saving the value to core state cache.
 /// The final value is known.
 pub(crate) struct DiceWorkerStateFinishedAndCached {
-    _prevent_cancellation: DisableCancellationGuard,
+    pub(crate) value: DiceComputedValue,
+    pub(crate) _prevent_cancellation: DisableCancellationGuard,
 }