event_bus.h

/* Copyright 2016 Google Inc. All Rights Reserved.
 
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
 
    http://www.apache.org/licenses/LICENSE-2.0
 
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
 
// TODO(b/25725560): Consider having a thread pool for invoking callbacks.
 
#ifndef TENSORFLOW_SERVING_UTIL_EVENT_BUS_H_
#define TENSORFLOW_SERVING_UTIL_EVENT_BUS_H_
 
#include <algorithm>
#include <functional>
#include <memory>
#include <type_traits>
#include <vector>
 
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/macros.h"
#include "tensorflow/core/platform/mutex.h"
#include "tensorflow/core/platform/thread_annotations.h"
#include "tensorflow/core/platform/types.h"
 
namespace tensorflow {
namespace serving {
 
template <typename E>
class EventBus : public std::enable_shared_from_this<EventBus<E>> {
  static_assert(std::is_move_assignable<E>::value, "E must be moveable");
 
 public:
  class Subscription {
   public:
    ~Subscription();
 
   private:
    friend class EventBus;
 
    explicit Subscription(std::weak_ptr<EventBus<E>> bus);
 
    // Weak pointer to the EventBus that originated this Subscription.
    std::weak_ptr<EventBus<E>> bus_;
 
    TF_DISALLOW_COPY_AND_ASSIGN(Subscription);
  };
 
  struct Options {
    // The environment to use for time.
    Env* env = Env::Default();
  };
 
  static std::shared_ptr<EventBus> CreateEventBus(const Options& options = {});
 
  ~EventBus() = default;
 
  struct EventAndTime {
    const E& event;
    uint64_t event_time_micros;
  };
 
  using Callback = std::function<void(const EventAndTime&)>;
 
  std::unique_ptr<Subscription> Subscribe(const Callback& callback)
      TF_LOCKS_EXCLUDED(mutex_) TF_MUST_USE_RESULT;
 
  void Publish(const E& event) TF_LOCKS_EXCLUDED(mutex_);
 
 private:
  explicit EventBus(const Options& options);
 
  // Unsubscribes the specified subscriber. Called only by Subscription.
  void Unsubscribe(const Subscription* subscription) TF_LOCKS_EXCLUDED(mutex_);
 
  // All of the information needed for a single subscription, both for
  // publishing events and unsubscribing.
  struct SubscriptionTuple {
    // Uniquely identifies the Subscription.
    Subscription* subscription;
    Callback callback;
  };
 
  // Mutex held for all operations on an EventBus including all publishing and
  // subscription operations.
  mutable mutex mutex_;
 
  // All subscriptions that the EventBus is aware of. Note that this is not
  // optimized for high scale in the number of subscribers.
  std::vector<SubscriptionTuple> subscriptions_ TF_GUARDED_BY(mutex_);
 
  const Options options_;
 
  TF_DISALLOW_COPY_AND_ASSIGN(EventBus);
};
 
// --- Implementation details below ---
 
template <typename E>
EventBus<E>::Subscription::Subscription(std::weak_ptr<EventBus<E>> bus)
    : bus_(std::move(bus)) {}
 
template <typename E>
EventBus<E>::Subscription::~Subscription() {
  std::shared_ptr<EventBus<E>> temp_shared_ptr = bus_.lock();
  if (temp_shared_ptr != nullptr) {
    temp_shared_ptr->Unsubscribe(this);
  }
}
 
template <typename E>
std::unique_ptr<typename EventBus<E>::Subscription> EventBus<E>::Subscribe(
    const Callback& callback) {
  mutex_lock lock(mutex_);
  std::unique_ptr<Subscription> subscription(
      new Subscription(this->shared_from_this()));
  subscriptions_.push_back({subscription.get(), callback});
  return subscription;
}
 
template <typename E>
EventBus<E>::EventBus(const Options& options) : options_(options) {}
 
template <typename E>
std::shared_ptr<EventBus<E>> EventBus<E>::CreateEventBus(
    const Options& options) {
  return std::shared_ptr<EventBus<E>>(new EventBus<E>(options));
}
 
template <typename E>
void EventBus<E>::Unsubscribe(
    const typename EventBus<E>::Subscription* subscription) {
  mutex_lock lock(mutex_);
  subscriptions_.erase(
      std::remove_if(subscriptions_.begin(), subscriptions_.end(),
                     [subscription](SubscriptionTuple s) {
                       return s.subscription == subscription;
                     }),
      subscriptions_.end());
}
 
template <typename E>
void EventBus<E>::Publish(const E& event) {
  mutex_lock lock(mutex_);
  const uint64_t event_time = options_.env->NowMicros();
  const EventAndTime event_and_time = {event, event_time};
  for (const SubscriptionTuple& subscription : subscriptions_) {
    subscription.callback(event_and_time);
  }
}
 
}  // namespace serving
}  // namespace tensorflow
 
#endif  // TENSORFLOW_SERVING_UTIL_EVENT_BUS_H_