xref: /aosp_15_r20/external/federated-compute/fcp/client/interruptible_runner.h (revision 14675a029014e728ec732f129a32e299b2da0601)
1 /*
2  * Copyright 2020 Google LLC
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *      http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 #ifndef FCP_CLIENT_INTERRUPTIBLE_RUNNER_H_
17 #define FCP_CLIENT_INTERRUPTIBLE_RUNNER_H_
18 
19 #include <functional>
20 #include <memory>
21 
22 #include "absl/status/status.h"
23 #include "absl/time/time.h"
24 #include "fcp/base/future.h"
25 #include "fcp/base/monitoring.h"
26 #include "fcp/base/scheduler.h"
27 #include "fcp/client/diag_codes.pb.h"
28 #include "fcp/client/log_manager.h"
29 
30 namespace fcp {
31 namespace client {
32 
33 // An executor that runs operations in a background thread, polling a callback
34 // periodically whether to abort, and aborting the operation if necessary.
35 // This uses a single-threaded thread pool. During execution of an operation,
36 // should_abort is polled periodically (polling_period), and if it returns true,
37 // the abort_function supplied along with the operation is called. The operation
38 // is then expected to abort within graceful_shutdown_period. If not, a diag
39 // code is logged and we wait for some time longer (extended_shutdown_period),
40 // and if the operation still does not finish, the program exits.
41 // The destructor blocks until the background thread has become idle.
42 class InterruptibleRunner {
43  public:
44   // A struct used to group polling & timeout related parameters.
45   struct TimingConfig {
46     absl::Duration polling_period;
47     absl::Duration graceful_shutdown_period;
48     absl::Duration extended_shutdown_period;
49   };
50 
51   // A struct used to group diagnostics related parameters.
52   struct DiagnosticsConfig {
53     ProdDiagCode interrupted;
54     ProdDiagCode interrupt_timeout;
55     ProdDiagCode interrupted_extended;
56     ProdDiagCode interrupt_timeout_extended;
57   };
58 
InterruptibleRunner(LogManager * log_manager,std::function<bool ()> should_abort,const TimingConfig & timing_config,const DiagnosticsConfig & diagnostics_config)59   InterruptibleRunner(LogManager* log_manager,
60                       std::function<bool()> should_abort,
61                       const TimingConfig& timing_config,
62                       const DiagnosticsConfig& diagnostics_config)
63       : log_manager_(log_manager),
64         should_abort_(should_abort),
65         timing_config_(timing_config),
66         diagnostics_config_(diagnostics_config) {
67     thread_pool_ = fcp::CreateThreadPoolScheduler(1);
68   }
69 
~InterruptibleRunner()70   ~InterruptibleRunner() { thread_pool_->WaitUntilIdle(); }
71 
72   // Executes f() on a background. Returns CANCELLED if the background thread
73   // was aborted, or a Status object from the background thread on successful
74   // completion.
75   absl::Status Run(std::function<absl::Status()> f,
76                    std::function<void()> abort_function);
77 
78  private:
79   absl::Status WaitUntilDone(fcp::thread::Future<absl::Status>&& run_future,
80                              std::function<void()> abort_function);
81   absl::Status Abort(fcp::thread::Future<absl::Status> run_future,
82                      std::function<void()> abort_function);
83 
84   std::unique_ptr<Scheduler> thread_pool_;
85   LogManager* const log_manager_;
86   std::function<bool()> should_abort_;
87   TimingConfig timing_config_;
88   DiagnosticsConfig diagnostics_config_;
89 };
90 
91 }  // namespace client
92 }  // namespace fcp
93 
94 #endif  // FCP_CLIENT_INTERRUPTIBLE_RUNNER_H_
95