tlbmc/hft/core/manager_fake.h

#ifndef THIRD_PARTY_MILOTIC_EXTERNAL_CC_TLBMC_HFT_CORE_MANAGER_FAKE_H_
#define THIRD_PARTY_MILOTIC_EXTERNAL_CC_TLBMC_HFT_CORE_MANAGER_FAKE_H_

#include <atomic>
#include <chrono>  // NOLINT
#include <cstdint>
#include <memory>
#include <thread>  // NOLINT
#include <utility>
#include <vector>

#include "absl/base/thread_annotations.h"
#include "absl/functional/any_invocable.h"
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/synchronization/mutex.h"
#include "absl/synchronization/notification.h"
#include "payload.pb.h"
#include "subscription_params.pb.h"
#include "tlbmc/hft/core/manager.h"

namespace milotic_hft {

// A fake implementation of the subscription manager.
// This class is thread-compatible.
// Call `SetFakeData` before calling `AddSubscription`.
class SubscriptionManagerFake : public SubscriptionManager {
 public:
  void SetFakeData(Payload&& payload) { payload_ = std::move(payload); }

  // The background thread will create fake data according to the subscription
  // params and call the callback.
  absl::StatusOr<std::shared_ptr<Subscription>> AddSubscription(
      const SubscriptionParams& params,
      absl::AnyInvocable<void(Payload&&)> on_data_callback) override {
    absl::MutexLock lock(mutex_);
    auto fake_subscription = std::make_shared<FakeSubscription>();
    active_fake_subscriptions_.push_back(fake_subscription);

    background_threads_.push_back(
        std::thread([params, on_data_callback = std::move(on_data_callback),
                     fake_subscription, this]() mutable {
          for (int32_t i = 0; i < params.num_batches(); ++i) {
            if (fake_subscription->IsCancelled()) {
              break;
            }
            // Sleep just for 1ms instead of `params.sampling_interval_ms()` to
            // speed up the test.
            std::this_thread::sleep_for(std::chrono::milliseconds(1));
            Payload payload = payload_;
            on_data_callback(std::move(payload));
          }
        }));
    return fake_subscription;
  }

  absl::Status Unsubscribe(
      const std::shared_ptr<Subscription>& subscription) override {
    absl::MutexLock lock(mutex_);
    auto fake_sub = std::dynamic_pointer_cast<FakeSubscription>(subscription);
    if (!fake_sub) {
      return absl::InvalidArgumentError(
          "Subscription is not a FakeSubscription");
    }

    // Signal cancellation to the background thread.
    fake_sub->Cancel();

    // Remove from active_fake_subscriptions_.
    for (auto it = active_fake_subscriptions_.begin();
         it != active_fake_subscriptions_.end(); ++it) {
      if (*it == fake_sub) {
        active_fake_subscriptions_.erase(it);
        break;
      }
    }
    return absl::OkStatus();
  }

  ~SubscriptionManagerFake() override {
    // Ensure all background threads are signaled to stop.
    absl::MutexLock lock(mutex_);
    for (const auto& sub : active_fake_subscriptions_) {
      sub->Cancel();
    }
    // Clear active_fake_subscriptions_ to release shared_ptrs.
    active_fake_subscriptions_.clear();

    // Join all threads. Threads should exit quickly after being cancelled.
    for (std::thread& thread : background_threads_) {
      if (thread.joinable()) {
        thread.join();
      }
    }
  }

 private:
  class FakeSubscription : public Subscription {
   public:
    // Signals the associated thread to stop.
    void Cancel() {
      if (!cancelled_.HasBeenNotified()) {
        cancelled_.Notify();
      }
    }

    // Returns true if the subscription has been cancelled.
    bool IsCancelled() const { return cancelled_.HasBeenNotified(); }

   private:
    absl::Notification cancelled_;
  };

  Payload payload_;
  absl::Mutex mutex_;
  // Keeps track of active FakeSubscription shared pointers.
  std::vector<std::shared_ptr<FakeSubscription>> active_fake_subscriptions_
      ABSL_GUARDED_BY(mutex_);
  // Stores the background threads. These are joined in the destructor.
  std::vector<std::thread> ABSL_GUARDED_BY(mutex_) background_threads_;
};

}  // namespace milotic_hft
#endif  // THIRD_PARTY_MILOTIC_EXTERNAL_CC_TLBMC_HFT_CORE_MANAGER_FAKE_H_