/* Copyright 2016 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.
==============================================================================*/

#include "tensorflow/core/distributed_runtime/local_master.h"

#include <unordered_map>

#include "tensorflow/core/distributed_runtime/master.h"
#include "tensorflow/core/platform/mutex.h"

namespace tensorflow {

namespace {
Status WaitForNotification(CallOptions* call_options,
                           const int64 default_timeout_in_ms, Notification* n) {
  int64 timeout_in_ms = call_options->GetTimeout();
  if (timeout_in_ms == 0) {
    timeout_in_ms = default_timeout_in_ms;
  }
  if (timeout_in_ms > 0) {
    int64 timeout_in_us = timeout_in_ms * 1000;
    bool notified = WaitForNotificationWithTimeout(n, timeout_in_us);
    if (!notified) {
      call_options->StartCancel();
      // The call has borrowed pointers to the request and response
      // messages, so we must still wait for the call to complete.
      n->WaitForNotification();
      return errors::DeadlineExceeded("Operation timed out.");
    }
  } else {
    n->WaitForNotification();
  }
  return Status::OK();
}
}  // namespace

LocalMaster::LocalMaster(Master* master_impl, const int64 default_timeout_in_ms)
    : master_impl_(master_impl),
      default_timeout_in_ms_(default_timeout_in_ms) {}

Status LocalMaster::CreateSession(CallOptions* call_options,
                                  const CreateSessionRequest* request,
                                  CreateSessionResponse* response) {
  Notification n;
  Status ret;
  master_impl_->CreateSession(request, response, [&n, &ret](const Status& s) {
    ret.Update(s);
    n.Notify();
  });
  TF_RETURN_IF_ERROR(
      WaitForNotification(call_options, default_timeout_in_ms_, &n));
  return ret;
}

Status LocalMaster::ExtendSession(CallOptions* call_options,
                                  const ExtendSessionRequest* request,
                                  ExtendSessionResponse* response) {
  Notification n;
  Status ret;
  master_impl_->ExtendSession(request, response, [&n, &ret](const Status& s) {
    ret.Update(s);
    n.Notify();
  });
  TF_RETURN_IF_ERROR(
      WaitForNotification(call_options, default_timeout_in_ms_, &n));
  return ret;
}

Status LocalMaster::PartialRunSetup(CallOptions* call_options,
                                    const PartialRunSetupRequest* request,
                                    PartialRunSetupResponse* response) {
  Notification n;
  Status ret;
  master_impl_->PartialRunSetup(request, response, [&n, &ret](const Status& s) {
    ret.Update(s);
    n.Notify();
  });
  TF_RETURN_IF_ERROR(
      WaitForNotification(call_options, default_timeout_in_ms_, &n));
  return ret;
}

Status LocalMaster::RunStep(CallOptions* call_options,
                            RunStepRequestWrapper* request,
                            MutableRunStepResponseWrapper* response) {
  Notification n;
  Status ret;
  master_impl_->RunStep(call_options, request, response,
                        [&n, &ret](const Status& s) {
                          ret.Update(s);
                          n.Notify();
                        });
  TF_RETURN_IF_ERROR(
      WaitForNotification(call_options, default_timeout_in_ms_, &n));
  return ret;
}

MutableRunStepRequestWrapper* LocalMaster::CreateRunStepRequest() {
  return new InMemoryRunStepRequest;
}

MutableRunStepResponseWrapper* LocalMaster::CreateRunStepResponse() {
  return new InMemoryRunStepResponse;
}

Status LocalMaster::CloseSession(CallOptions* call_options,
                                 const CloseSessionRequest* request,
                                 CloseSessionResponse* response) {
  Notification n;
  Status ret;
  master_impl_->CloseSession(request, response, [&n, &ret](const Status& s) {
    ret.Update(s);
    n.Notify();
  });
  TF_RETURN_IF_ERROR(
      WaitForNotification(call_options, default_timeout_in_ms_, &n));
  return ret;
}

Status LocalMaster::ListDevices(CallOptions* call_options,
                                const ListDevicesRequest* request,
                                ListDevicesResponse* response) {
  Notification n;
  Status ret;
  master_impl_->ListDevices(request, response, [&n, &ret](const Status& s) {
    ret.Update(s);
    n.Notify();
  });
  TF_RETURN_IF_ERROR(
      WaitForNotification(call_options, default_timeout_in_ms_, &n));
  return ret;
}

Status LocalMaster::Reset(CallOptions* call_options,
                          const ResetRequest* request,
                          ResetResponse* response) {
  Notification n;
  Status ret;
  master_impl_->Reset(request, response, [&n, &ret](const Status& s) {
    ret.Update(s);
    n.Notify();
  });
  TF_RETURN_IF_ERROR(
      WaitForNotification(call_options, default_timeout_in_ms_, &n));
  return ret;
}

Status LocalMaster::MakeCallable(CallOptions* call_options,
                                 const MakeCallableRequest* request,
                                 MakeCallableResponse* response) {
  Notification n;
  Status ret;
  master_impl_->MakeCallable(request, response, [&n, &ret](const Status& s) {
    ret.Update(s);
    n.Notify();
  });
  TF_RETURN_IF_ERROR(
      WaitForNotification(call_options, default_timeout_in_ms_, &n));
  return ret;
}
Status LocalMaster::RunCallable(CallOptions* call_options,
                                const RunCallableRequest* request,
                                RunCallableResponse* response) {
  Notification n;
  Status ret;
  master_impl_->RunCallable(call_options, request, response,
                            [&n, &ret](const Status& s) {
                              ret.Update(s);
                              n.Notify();
                            });
  TF_RETURN_IF_ERROR(
      WaitForNotification(call_options, default_timeout_in_ms_, &n));
  return ret;
}
Status LocalMaster::ReleaseCallable(CallOptions* call_options,
                                    const ReleaseCallableRequest* request,
                                    ReleaseCallableResponse* response) {
  Notification n;
  Status ret;
  master_impl_->ReleaseCallable(request, response, [&n, &ret](const Status& s) {
    ret.Update(s);
    n.Notify();
  });
  TF_RETURN_IF_ERROR(
      WaitForNotification(call_options, default_timeout_in_ms_, &n));
  return ret;
}

namespace {
mutex* get_local_master_registry_lock() {
  static mutex local_master_registry_lock(LINKER_INITIALIZED);
  return &local_master_registry_lock;
}

struct MasterInfo {
  Master* master;
  const int64 default_timeout_in_ms;

  MasterInfo(Master* master, const int64 default_timeout_in_ms)
      : master(master), default_timeout_in_ms(default_timeout_in_ms) {}
};

typedef std::unordered_map<string, MasterInfo> LocalMasterRegistry;
LocalMasterRegistry* local_master_registry() {
  static LocalMasterRegistry* local_master_registry_ = new LocalMasterRegistry;
  return local_master_registry_;
}
}  // namespace

/* static */
void LocalMaster::Register(const string& target, Master* master,
                           int64 default_timeout_in_ms) {
  mutex_lock l(*get_local_master_registry_lock());
  local_master_registry()->insert(
      {target, MasterInfo(master, default_timeout_in_ms)});
}

/* static */
std::unique_ptr<LocalMaster> LocalMaster::Lookup(const string& target) {
  std::unique_ptr<LocalMaster> ret;
  mutex_lock l(*get_local_master_registry_lock());
  auto iter = local_master_registry()->find(target);
  if (iter != local_master_registry()->end()) {
    ret.reset(new LocalMaster(iter->second.master,
                              iter->second.default_timeout_in_ms));
  }
  return ret;
}

}  // namespace tensorflow