ml/process.h

// Copyright 2021 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#ifndef ML_PROCESS_H_
#define ML_PROCESS_H_

#include <memory>
#include <string>
#include <unordered_map>

#include <unistd.h>

#include <base/callback_forward.h>
#include <base/macros.h>
#include <base/no_destructor.h>
#include <base/process/process_metrics.h>
#include <mojo/public/cpp/bindings/remote.h>
#include <base/sequence_checker.h>

#include "ml/machine_learning_service_impl.h"

namespace ml {

// A singleton class to store the global process information and provide
// process management functions.
// Usage: access the global instance by calling `Process::GetInstance()`.
class Process {
 public:
  // The type of a process.
  // "kControlForTest" denotes the control process in unit tests (i.e. the
  // process that runs the ml_service_test binary). "kSingleProcessForTest"
  // means the program will not spawn worker processes and use one single
  // process for testing.
  enum class Type {
    kUnset = 0,
    kControl = 1,
    kWorker = 2,
    kControlForTest = 3,        // Like control process but with less strict
                                // sandboxing for using in testing.
    kSingleProcessForTest = 4,  // Temporary, used by old single-process tests.
  };

  // The exit code of a process.
  enum ExitCode : int {
    kSuccess = 0,
    // Only for worker process used when its mojo connection with the control
    // process breaks.
    kWorkerDisconnectWithControl = 1,
    kInvalidProcessType = 2,
    kUnexpectedCommandLine = 3,
  };

  // The worker process info, containing object to contact and measure worker
  // process in the control process.
  struct WorkerInfo {
    // The Mojo remote to call the worker process's `MachineLearningService`
    // bindings.
    mojo::Remote<chromeos::machine_learning::mojom::MachineLearningService>
        remote;
    // The process metrics object of the worker process.
    std::unique_ptr<base::ProcessMetrics> process_metrics;
  };

  static Process* GetInstance();

  int Run(int argc, char* argv[]);

  // Gets the process type of current process.
  Type GetType();

  // Returns true if the worker process has been started successfully and the
  // worker's pid is stored in `worker_pid`. Otherwise returns false and
  // `worker_pid` is unchanged.
  // The argument `model_name` has two usages:
  //   - it used in logging (like `metrics_model_name`).
  //   - it also determines which seccomp policy list to use in sandboxing the
  //     worker process.
  bool SpawnWorkerProcessAndGetPid(const mojo::PlatformChannel& channel,
                                   const std::string& model_name,
                                   pid_t* worker_pid);

  // Returns a reference of the remote of the worker process. The remote is hold
  // in the `worker_pid_info_map_` object.
  mojo::Remote<chromeos::machine_learning::mojom::MachineLearningService>&
  SendMojoInvitationAndGetRemote(pid_t worker_pid,
                                 mojo::PlatformChannel channel,
                                 const std::string& model_name);

  // Removes a worker process from metadata. This does not terminate the
  // worker process.
  void UnregisterWorkerProcess(pid_t pid);

  const std::unordered_map<pid_t, WorkerInfo>& GetWorkerPidInfoMap();

  // Sets the process type. Only used in testing.
  void SetTypeForTesting(Type type);

  // Sets the path of mlservice. Only used in testing.
  void SetMlServicePathForTesting(const std::string& path);

  // Sets the `before_exit_worker_disconnect_handler_hook`, only used in
  // testing.
  void SetBeforeExitWorkerDisconnectHandlerHookForTesting(
      base::RepeatingClosure hook);

  // Returns if the current process is a control process (i.e. `kControl ||
  // kControlForTest`).
  bool IsControlProcess();

  // Returns if the current process is a worker process (i.e. that will do the
  // actually inference work, `kWorker || kSingleProcessForTest`).
  bool IsWorkerProcess();

 private:
  friend base::NoDestructor<Process>;

  Process();
  ~Process();

  // Can only be called by the control process.
  void ControlProcessRun();

  // Can only be called by the worker process.
  // Input: the file descriptor used to bootstrap Mojo connection.
  void WorkerProcessRun();

  // The disconnect handler of control process for the mojo connection to the
  // worker process.
  void InternalPrimordialMojoPipeDisconnectHandler(pid_t child_pid);

  // The type of current process.
  Type process_type_;

  // The file descriptor to bootstrap the mojo connection of current process.
  // Only meaningful for worker process.
  int mojo_bootstrap_fd_;

  // Path to the ml_service binary. Normally (and by default) it is
  // "/usr/bin/ml_service". We may change the value here for testing.
  std::string ml_service_path_;

  // The map from pid to the info of worker processes. Only meaningful for
  // control process.
  std::unordered_map<pid_t, WorkerInfo> worker_pid_info_map_;

  // The function called in the `kControlForTesting` process at the last of
  // disconnection handler of the mojo connection to the worker process, only
  // used in testing.
  base::RepeatingClosure before_exit_worker_disconnect_handler_hook_;

  // Mainly used for guarding `worker_pid_info_map_`.
  SEQUENCE_CHECKER(sequence_checker_);
};

}  // namespace ml

#endif  // ML_PROCESS_H_