// Licensed to the Apache Software Foundation (ASF) under one

// or more contributor license agreements.  See the NOTICE file

// distributed with this work for additional information

// regarding copyright ownership.  The ASF licenses this file

// to you 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.

// Most code of this file is copied from rocksdb SyncPoint.

// https://github.com/facebook/rocksdb

// clang-format off

#include "sync_point.h"

#include <atomic>

#include <condition_variable>

#include <functional>

#include <mutex>

#include <random>

#include <string>

#include <thread>

#include <unordered_map>

#include <unordered_set>

namespace doris {

struct SyncPoint::Data { // impl

public:

  Data() : enabled_(false) { }

  virtual ~Data() {}

  void process(const std::string& point, std::vector<std::any>&& cb_args);

  void load_dependency(const std::vector<SyncPointPair>& dependencies);

  void load_dependency_and_markers(

                                const std::vector<SyncPointPair>& dependencies,

                                const std::vector<SyncPointPair>& markers);

  bool predecessors_all_cleared(const std::string& point);

  void set_call_back(const std::string& point,

                    const std::function<void(std::vector<std::any>&&)>& callback);

  void clear_call_back(const std::string& point);

  void clear_all_call_backs();

  void enable_processing();

  void disable_processing();

  void clear_trace();

private:

  bool disable_by_marker(const std::string& point, std::thread::id thread_id);

private:

  // successor/predecessor map loaded from load_dependency

  std::unordered_map<std::string, std::vector<std::string>> successors_;

  std::unordered_map<std::string, std::vector<std::string>> predecessors_;

  std::unordered_map<std::string, std::function<void(std::vector<std::any>&&)>> callbacks_;

  std::unordered_map<std::string, std::vector<std::string>> markers_;

  std::unordered_map<std::string, std::thread::id> marked_thread_id_;

  std::mutex mutex_;

  std::condition_variable cv_;

  // sync points that have been passed through

  std::unordered_set<std::string> cleared_points_;

  std::atomic<bool> enabled_;

  int num_callbacks_running_ = 0;

};

SyncPoint* SyncPoint::get_instance() {

  static SyncPoint sync_point;

  return &sync_point;

}

SyncPoint::SyncPoint() : 

  impl_(new Data) {

}

SyncPoint:: ~SyncPoint() {

  delete impl_;

}

void SyncPoint::load_dependency(const std::vector<SyncPointPair>& dependencies) {

  impl_->load_dependency(dependencies);

}

void SyncPoint::load_dependency_and_markers(

                                const std::vector<SyncPointPair>& dependencies,

                                const std::vector<SyncPointPair>& markers) {

  impl_->load_dependency_and_markers(dependencies, markers);

}

void SyncPoint::set_call_back(const std::string& point,

                              const std::function<void(std::vector<std::any>&&)>& callback) {

  impl_->set_call_back(point, callback);

}

void SyncPoint::clear_call_back(const std::string& point) {

  impl_->clear_call_back(point);

}

void SyncPoint::clear_all_call_backs() {

  impl_->clear_all_call_backs();

}

void SyncPoint::enable_processing() {

  impl_->enable_processing();

}

void SyncPoint::disable_processing() {

  impl_->disable_processing();

}

void SyncPoint::clear_trace() {

  impl_->clear_trace();

}

void SyncPoint::process(const std::string& point, std::vector<std::any>&& cb_arg) {

  impl_->process(point, std::move(cb_arg));

}

// =============================================================================

// SyncPoint implementation

// =============================================================================

void SyncPoint::Data::load_dependency(

                               const std::vector<SyncPointPair>& dependencies) {

  std::lock_guard lock(mutex_);

  successors_.clear();

  predecessors_.clear();

  cleared_points_.clear();

  for (const auto& dependency : dependencies) {

    successors_[dependency.predecessor].push_back(dependency.successor);

    predecessors_[dependency.successor].push_back(dependency.predecessor);

  }

  cv_.notify_all();

}

/**

 * Markers are also dependency descriptions

 */

void SyncPoint::Data::load_dependency_and_markers(

                                const std::vector<SyncPointPair>& dependencies,

                                const std::vector<SyncPointPair>& markers) {

  std::lock_guard lock(mutex_);

  successors_.clear();

  predecessors_.clear();

  cleared_points_.clear();

  markers_.clear();

  marked_thread_id_.clear();

  for (const auto& dependency : dependencies) {

    successors_[dependency.predecessor].push_back(dependency.successor);

    predecessors_[dependency.successor].push_back(dependency.predecessor);

  }

  for (const auto& marker : markers) {

    successors_[marker.predecessor].push_back(marker.successor);

    predecessors_[marker.successor].push_back(marker.predecessor);

    markers_[marker.predecessor].push_back(marker.successor);

  }

  cv_.notify_all();

}

bool SyncPoint::Data::predecessors_all_cleared(const std::string& point) {

  for (const auto& pred : predecessors_[point]) {

    if (cleared_points_.count(pred) == 0) {

      return false;

    }

  }

  return true;

}

void SyncPoint::Data::clear_call_back(const std::string& point) {

  std::unique_lock lock(mutex_);

  callbacks_.erase(point);

}

void SyncPoint::Data::clear_all_call_backs() {

  std::unique_lock lock(mutex_);

  callbacks_.clear();

}

void SyncPoint::Data::process(const std::string& point, std::vector<std::any>&& cb_arg) {

  if (!enabled_) {

    return;

  }

  std::unique_lock lock(mutex_);

  auto thread_id = std::this_thread::get_id();

  auto marker_iter = markers_.find(point);

  // if current sync point is a marker

  // record it in marked_thread_id_ for all its successors

  if (marker_iter != markers_.end()) {

    for (auto& marked_point : marker_iter->second) {

      marked_thread_id_.emplace(marked_point, thread_id);

    }

  }

  // if current point is a marker's successor 

  if (disable_by_marker(point, thread_id)) {

    return;

  }

  while (!predecessors_all_cleared(point)) {

    cv_.wait(lock);

    if (disable_by_marker(point, thread_id)) {

      return;

    }

  }

  auto callback_pair = callbacks_.find(point);

  if (callback_pair != callbacks_.end()) {

    num_callbacks_running_++;

    auto callback = callback_pair->second; 

    mutex_.unlock();

    callback(std::move(cb_arg));

    mutex_.lock();

    num_callbacks_running_--;

  }

  cleared_points_.insert(point);

  cv_.notify_all();

}

bool SyncPoint::Data::disable_by_marker(const std::string& point,

                                        std::thread::id thread_id) {

  auto marked_point_iter = marked_thread_id_.find(point);

  return marked_point_iter != marked_thread_id_.end() // is a successor

          && thread_id != marked_point_iter->second;

}

void SyncPoint::Data::set_call_back(const std::string& point,

                                  const std::function<void(std::vector<std::any>&&)>& callback) {

  std::lock_guard lock(mutex_);

  callbacks_[point] = callback;

}

void SyncPoint::Data::clear_trace() {

  std::lock_guard lock(mutex_);

  cleared_points_.clear();

}

void SyncPoint::Data::enable_processing() {

  enabled_ = true;

}

void SyncPoint::Data::disable_processing() {

  enabled_ = false;

}

} // namespace doris

// clang-format on

// vim: et tw=80 ts=2 sw=2 cc=80: