/* Copyright 2019 The TensorFlow Authors. 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.
==============================================================================*/

#ifndef TENSORFLOW_CORE_DISTRIBUTED_RUNTIME_EAGER_REMOTE_MGR_H_
#define TENSORFLOW_CORE_DISTRIBUTED_RUNTIME_EAGER_REMOTE_MGR_H_

#include <unordered_map>
#include <vector>

#include "absl/strings/string_view.h"
#include "tensorflow/core/common_runtime/eager/eager_executor.h"
#include "tensorflow/core/common_runtime/eager/tensor_handle.h"
#include "tensorflow/core/distributed_runtime/eager/remote_tensor_handle.h"
#include "tensorflow/core/platform/mutex.h"

namespace tensorflow {
namespace eager {

// This class manages the states required to setup an eager cluster.
// TODO(fishx): Move remote state from context to this class.
class RemoteMgr {
 public:
  RemoteMgr(bool is_master, EagerContext* ctx)
      : is_master_(is_master), parent_(ctx) {}

  ~RemoteMgr() {
    for (const auto& entry : remote_tensor_handle_map_) {
      entry.second->Unref();
    }
  }

  bool IsMaster() { return is_master_; }

  void AddOperationOutputs(
      const absl::Span<tensorflow::TensorHandle* const> handles,
      int64_t operation_id);

  void AddOperationOutput(tensorflow::TensorHandle* handles,
                          int64_t operation_id, int32_t output_num);

  absl::Status GetTensorHandle(const RemoteTensorHandleInternal& remote_handle,
                               tensorflow::TensorHandle** handle);

  absl::Status DeleteTensorHandle(
      const RemoteTensorHandleInternal& remote_handle);

  // Helper function to create monotonically increasing ids unique to this
  // context.
  uint64_t NextOpId() {
    DCHECK(is_master_);
    mutex_lock l(next_id_mutex_);
    return next_op_id_++;
  }

  // Serialize a remote TensorHandle to a RemoteTensorHandle.
  // If wait_until_ready is true, block until the remote handle is ready on a
  // remote worker.
  absl::Status SerializeRemoteTensorHandle(
      TensorHandle* in, const bool wait_until_ready, RemoteTensorHandle* out,
      Device* device, absl::string_view device_name = "",
      const bool serialize_resource_dtype_and_shape = false);

  // Deserialize a RemoteTensorHandle to a TensorHandle(local/remote).
  // The output holds a reference to the TensorHandle.
  absl::Status DeserializeRemoteTensorHandle(const RemoteTensorHandle& in,
                                             TensorHandle** out);

  EagerExecutor& GetOrCreateExecutorForStream(uint64_t stream_id);

  void DeleteExecutorForStream(uint64_t stream_id);

 protected:
  mutex next_id_mutex_;
  uint64_t next_op_id_ TF_GUARDED_BY(next_id_mutex_) = 1;

 private:
  // Returns the op_id and output_num if the given local TensorHandle exists in
  // remote_tensor_handle_map_.
  absl::Status GetRemoteTensorHandle(const tensorflow::TensorHandle* handle,
                                     const bool wait_until_ready,
                                     int64_t* op_id, int32_t* output_num)
      TF_SHARED_LOCKS_REQUIRED(remote_tensor_handle_mu_);

  absl::Status GetTensorHandleImpl(
      const RemoteTensorHandleInternal& remote_handle,
      tensorflow::TensorHandle** handle)
      TF_SHARED_LOCKS_REQUIRED(remote_tensor_handle_mu_);

  absl::Status GetMirroredResourceShape(
      const RemoteTensorHandleInternal& remote_handle,
      std::vector<DtypeAndPartialTensorShape>* handle);

  bool is_master_;

  using RemoteTensorHandleMap =
      gtl::FlatMap<RemoteTensorHandleInternal, tensorflow::TensorHandle*,
                   RemoteTensorHandleInternalHash,
                   RemoteTensorHandleInternalEquals>;
  using MirroredResourceShapeMap = gtl::FlatMap<
      RemoteTensorHandleInternal, std::vector<DtypeAndPartialTensorShape>,
      RemoteTensorHandleInternalHash, RemoteTensorHandleInternalEquals>;

  mutex remote_tensor_handle_mu_;
  // This map maintains the TensorHandles that are required by remote workers
  // in the cluster. Each map key is generated by the master, so it should be
  // globally unique. This map owns references on the handles it contains.
  RemoteTensorHandleMap remote_tensor_handle_map_
      TF_GUARDED_BY(remote_tensor_handle_mu_);

  mutex mirrored_resource_shape_mu_;
  // This map maintains the data types and shapes of resource variables required
  // by remote workers in the cluster. Each map key is generated by the master,
  // so it should be globally unique.
  MirroredResourceShapeMap mirrored_resource_shape_map_
      TF_GUARDED_BY(mirrored_resource_shape_mu_);

  EagerContext* parent_;  // not owned.

  mutex executor_map_mu_;
  std::unordered_map<uint64_t, EagerExecutor> executor_map_
      TF_GUARDED_BY(executor_map_mu_);
};

}  // namespace eager
}  // namespace tensorflow

#endif  // TENSORFLOW_CORE_DISTRIBUTED_RUNTIME_EAGER_REMOTE_MGR_H_
