public/lib/callback/waiter.h

// Copyright 2018 The Fuchsia 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 LIB_CALLBACK_WAITER_H_
#define LIB_CALLBACK_WAITER_H_

#include <memory>
#include <utility>
#include <vector>

#include <lib/fit/function.h>

#include "lib/fxl/macros.h"
#include "lib/fxl/memory/ref_counted.h"

namespace callback {

namespace internal {

template <typename S, typename T>
class ResultAccumulator {
 public:
  explicit ResultAccumulator(S success_status)
      : success_status_(success_status), result_status_(success_status_) {}

  size_t PrepareCall() {
    results_.emplace_back();
    return results_.size() - 1;
  }

  bool Update(size_t index, S status, T result) {
    if (status != success_status_) {
      result_status_ = status;
      results_.clear();
      return false;
    }
    results_[index] = std::move(result);
    return true;
  }

  std::pair<S, std::vector<T>> Result() {
    return std::make_pair(result_status_, std::move(results_));
  }

 private:
  size_t current_index_ = 0;
  std::vector<T> results_;
  S success_status_;
  S result_status_;
};

template <typename S>
class StatusAccumulator {
 public:
  explicit StatusAccumulator(S success_status)
      : success_status_(success_status), result_status_(success_status_) {}

  bool PrepareCall() { return true; }

  bool Update(bool /*token*/, S status) {
    result_status_ = status;
    return success_status_ == result_status_;
  }

  S Result() { return result_status_; }

 private:
  S success_status_;
  S result_status_;
};

// AnyAccumulator is the accumulator for the AnyWaiter class.
// It continues until an |Update| call matches |success_status|.
template <typename S, typename V>
class AnyAccumulator {
 public:
  AnyAccumulator(S success_status, S default_status, V default_value)
      : success_status_(success_status),
        result_status_(default_status),
        value_(std::move(default_value)) {}

  bool PrepareCall() { return true; }

  bool Update(bool /*token*/, S status, V value) {
    if (status == success_status_) {
      value_ = std::move(value);
    }
    result_status_ = std::move(status);
    // Continue until we get a success.
    return result_status_ != success_status_;
  }

  std::pair<S, V> Result() {
    return std::make_pair(result_status_, std::move(value_));
  }

 private:
  const S success_status_;
  S result_status_;
  V value_;
};

template <typename S, typename V>
class PromiseAccumulator {
 public:
  PromiseAccumulator(S default_status, V default_value)
      : status_(default_status), value_(std::move(default_value)) {}

  bool PrepareCall() { return true; }

  bool Update(bool /*token*/, S status, V value) {
    status_ = std::move(status);
    value_ = std::move(value);
    return false;
  }

  std::pair<S, V> Result() {
    return std::make_pair(status_, std::move(value_));
  }

 private:
  S status_;
  V value_;
};

class CompletionAccumulator {
 public:
  bool PrepareCall() { return true; }

  bool Update(bool /*token*/) { return true; }

  bool Result() { return true; }
};

}  // namespace internal

// Base implementation for all specialized waiters. It uses an Accumulator
// abstraction to aggregate results from the different callbacks.
// |A| is the accumulator, |R| is the final return type and |Args| are the
// arguments of the callbacks. The accumulator must have the following
// interface:
// class Accumulator {
//  public:
//   TOKEN PrepareCall();
//   bool Update(TOKEN token, Args... args);
//   R Result();
// };
template <typename A, typename R, typename... Args>
class BaseWaiter : public fxl::RefCountedThreadSafe<BaseWaiter<A, R, Args...>> {
 public:
  fit::function<void(Args...)> NewCallback() {
    FXL_DCHECK(!finalized_) << "Waiter was already finalized.";
    FXL_DCHECK(!cancelled_) << "Waiter has been cancelled.";
    if (done_) {
      return [](Args...) {};
    }
    ++pending_callbacks_;
    return [waiter_ref = fxl::RefPtr<BaseWaiter<A, R, Args...>>(this),
            token = accumulator_.PrepareCall()](Args&&... args) mutable {
      waiter_ref->ReturnResult(std::move(token), std::forward<Args>(args)...);
    };
  }

  void Finalize(fit::function<void(R)> callback) {
    FXL_DCHECK(!finalized_) << "Waiter already finalized, can't finalize more!";
    FXL_DCHECK(!cancelled_) << "Waiter has been cancelled.";
    result_callback_ = std::move(callback);
    finalized_ = true;
    ExecuteCallbackIfFinished();
  }

  void Cancel() { cancelled_ = true; }

 protected:
  explicit BaseWaiter(A&& accumulator) : accumulator_(std::move(accumulator)) {}
  virtual ~BaseWaiter() {}

 private:
  FRIEND_REF_COUNTED_THREAD_SAFE(BaseWaiter);
  FRIEND_MAKE_REF_COUNTED(BaseWaiter);

  template <typename T>
  void ReturnResult(T token, Args... args) {
    if (done_) {
      FXL_DCHECK(!pending_callbacks_);
      return;
    }
    done_ = !accumulator_.Update(std::move(token), std::forward<Args>(args)...);
    if (done_) {
      pending_callbacks_ = 0;
    } else {
      --pending_callbacks_;
    }
    ExecuteCallbackIfFinished();
  }

  void ExecuteCallbackIfFinished() {
    FXL_DCHECK(!finished_) << "Waiter already finished.";
    if (!finalized_ || pending_callbacks_) {
      return;
    }
    finished_ = true;
    if (!cancelled_) {
      result_callback_(accumulator_.Result());
      // The callback might delete this class.
      return;
    }
  }

  A accumulator_;
  // TODO(LE-382): Simplify the state machine here.
  bool done_ = false;
  bool finalized_ = false;
  bool finished_ = false;
  bool cancelled_ = false;
  size_t pending_callbacks_ = 0;

  fit::function<void(R)> result_callback_;
};

// Waiter can be used to collate the results of many asynchronous calls into one
// callback. A typical usage example would be:
// auto waiter = fxl::MakeRefCounted<callback::Waiter<Status,
//                                   std::unique_ptr<Object>>>(Status::OK);
// storage->GetObject(object_digest1, waiter->NewCallback());
// storage->GetObject(object_digest2, waiter->NewCallback());
// storage->GetObject(object_digest3, waiter->NewCallback());
// ...
// waiter->Finalize([](Status s, std::vector<std::unique_ptr<Object>> v) {
//   do something with the returned objects
// });
template <class S, class T>
class Waiter : public BaseWaiter<internal::ResultAccumulator<S, T>,
                                 std::pair<S, std::vector<T>>, S, T> {
 public:
  void Finalize(fit::function<void(S, std::vector<T>)> callback) {
    BaseWaiter<internal::ResultAccumulator<S, T>, std::pair<S, std::vector<T>>,
               S, T>::Finalize([callback =
                                    std::move(callback)](
                                   std::pair<S, std::vector<T>> result) {
      callback(result.first, std::move(result.second));
    });
  }

 private:
  FRIEND_REF_COUNTED_THREAD_SAFE(Waiter);
  FRIEND_MAKE_REF_COUNTED(Waiter);
  ~Waiter() override{};

  explicit Waiter(S success_status)
      : BaseWaiter<internal::ResultAccumulator<S, T>,
                   std::pair<S, std::vector<T>>, S, T>(
            internal::ResultAccumulator<S, T>(success_status)) {}
};

// StatusWaiter can be used to collate the results of many asynchronous calls
// into one callback. It is different from Waiter in that the callbacks only use
// S (e.g. storage::Status) as an argument.
template <class S>
class StatusWaiter : public BaseWaiter<internal::StatusAccumulator<S>, S, S> {
 private:
  FRIEND_REF_COUNTED_THREAD_SAFE(StatusWaiter);
  FRIEND_MAKE_REF_COUNTED(StatusWaiter);

  explicit StatusWaiter(S success_status)
      : BaseWaiter<internal::StatusAccumulator<S>, S, S>(
            internal::StatusAccumulator<S>(success_status)) {}
  ~StatusWaiter() override{};
};

// AnyWaiter is used to wait many asynchronous calls and returns the first
// successful result. It will return |default_status| and |default_value| only
// if no callback was called with a |success_status| status.
template <class S, class V>
class AnyWaiter
    : public BaseWaiter<internal::AnyAccumulator<S, V>, std::pair<S, V>, S, V> {
 public:
  void Finalize(fit::function<void(S, V)> callback) {
    BaseWaiter<internal::AnyAccumulator<S, V>, std::pair<S, V>, S, V>::Finalize(
        [callback = std::move(callback)](std::pair<S, V> result) {
          callback(result.first, std::move(result.second));
        });
  }

 private:
  FRIEND_REF_COUNTED_THREAD_SAFE(AnyWaiter);
  FRIEND_MAKE_REF_COUNTED(AnyWaiter);

  // Creates a new waiter. |success_status| and |default_value| will be
  // returned to the callback in |Finalize| if |NewCallback| is not called.
  AnyWaiter(S success_status, S default_status, V default_value = V())
      : BaseWaiter<internal::AnyAccumulator<S, V>, std::pair<S, V>, S, V>(
            internal::AnyAccumulator<S, V>(success_status, default_status,
                                           std::move(default_value))) {}
  ~AnyWaiter() override{};
};

// Promise is used to wait on a single asynchronous call. A typical usage
// example is:
// auto promise =
//     fxl::MakeRefCounted<callback::Promise<Status, std::unique_ptr<Object>>>(
//         Status::ILLEGAL_STATE);
// storage->GetObject(object_digest1, promise->NewCallback());
// ...
//
// promise->Finalize([](Status s, std::unique_ptr<Object> o) {
//   do something with the returned object
// });
template <class S, class V>
class Promise : public BaseWaiter<internal::PromiseAccumulator<S, V>,
                                  std::pair<S, V>, S, V> {
 public:
  void Finalize(fit::function<void(S, V)> callback) {
    BaseWaiter<internal::PromiseAccumulator<S, V>, std::pair<S, V>, S,
               V>::Finalize([callback =
                                 std::move(callback)](std::pair<S, V> result) {
      callback(result.first, std::move(result.second));
    });
  }

 private:
  FRIEND_REF_COUNTED_THREAD_SAFE(Promise);
  FRIEND_MAKE_REF_COUNTED(Promise);

  // Creates a new promise. |default_status| and |default_value| will be
  // returned to the callback in |Finalize| if |NewCallback| is not called.
  Promise(S default_status, V default_value = V())
      : BaseWaiter<internal::PromiseAccumulator<S, V>, std::pair<S, V>, S, V>(
            internal::PromiseAccumulator<S, V>(default_status,
                                               std::move(default_value))) {}
  ~Promise() override{};
};

// CompletionWaiter can be used to be notified on completion of a computation.
class CompletionWaiter
    : public BaseWaiter<internal::CompletionAccumulator, bool> {
 public:
  void Finalize(fit::function<void()> callback) {
    BaseWaiter<internal::CompletionAccumulator, bool>::Finalize(
        [callback = std::move(callback)](bool result) { callback(); });
  }

 private:
  FRIEND_REF_COUNTED_THREAD_SAFE(CompletionWaiter);
  FRIEND_MAKE_REF_COUNTED(CompletionWaiter);

  CompletionWaiter()
      : BaseWaiter<internal::CompletionAccumulator, bool>(
            internal::CompletionAccumulator()) {}
  ~CompletionWaiter() override{};
};

}  // namespace callback

#endif  // LIB_CALLBACK_WAITER_H_