/* Copyright 2021 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/data/service/worker_client.h"

#include <memory>
#include <string>
#include <utility>

#include "absl/memory/memory.h"
#include "absl/strings/substitute.h"
#include "absl/types/optional.h"
#include "tensorflow/core/data/service/common.pb.h"
#include "tensorflow/core/data/service/data_service.h"
#include "tensorflow/core/data/service/data_transfer.h"
#include "tensorflow/core/data/service/dispatcher.pb.h"
#include "tensorflow/core/data/service/dispatcher_client.h"
#include "tensorflow/core/data/service/test_cluster.h"
#include "tensorflow/core/data/service/test_util.h"
#include "tensorflow/core/data/service/worker.pb.h"
#include "tensorflow/core/data/service/worker_impl.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/tensor_testutil.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/platform/status_matchers.h"
#include "tensorflow/core/platform/statusor.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/platform/types.h"
#include "tensorflow/core/protobuf/error_codes.pb.h"

namespace tensorflow {
namespace data {
namespace {

using ::testing::MatchesRegex;

constexpr const char kProtocol[] = "grpc";

class WorkerClientTest : public ::testing::Test {
 protected:
  void SetUp() override {
    test_cluster_ = absl::make_unique<TestCluster>(/*num_workers=*/1);
    TF_ASSERT_OK(test_cluster_->Initialize());
    dispatcher_client_ = absl::make_unique<DataServiceDispatcherClient>(
        test_cluster_->DispatcherAddress(), kProtocol);
  }

  // Creates a dataset and returns the dataset ID.
  StatusOr<int64> RegisterDataset(const int64 range) {
    TF_ASSIGN_OR_RETURN(test_util::GraphDefTestCase graph_def_test_case,
                        test_util::map_test_case(range));
    int64 dataset_id = 0;
    TF_RETURN_IF_ERROR(dispatcher_client_->RegisterDataset(
        graph_def_test_case.graph_def, dataset_id));
    return dataset_id;
  }

  // Creates a job and returns the job client ID.
  StatusOr<int64> CreateJob(const int64 dataset_id) {
    int64 job_client_id = 0;
    TF_RETURN_IF_ERROR(dispatcher_client_->GetOrCreateJob(
        dataset_id, ProcessingMode::PARALLEL_EPOCHS, /*job_key=*/absl::nullopt,
        /*num_consumers=*/absl::nullopt, job_client_id));
    return job_client_id;
  }

  // Gets the task for job `job_client_id`.
  StatusOr<int64> GetTaskToRead(const int64 job_client_id) {
    ClientHeartbeatRequest request;
    ClientHeartbeatResponse response;
    request.set_job_client_id(job_client_id);
    TF_RETURN_IF_ERROR(dispatcher_client_->ClientHeartbeat(request, response));
    if (response.task_info().empty()) {
      return errors::NotFound(
          absl::Substitute("No task found for job $0.", job_client_id));
    }
    return response.task_info(0).task_id();
  }

  StatusOr<std::unique_ptr<DataServiceWorkerClient>> GetWorkerClient(
      const std::string& data_transfer_protocol) {
    return CreateDataServiceWorkerClient(
        GetWorkerAddress(), /*protocol=*/kProtocol, data_transfer_protocol);
  }

  StatusOr<GetElementResult> GetElement(DataServiceWorkerClient& client,
                                        const int64 task_id) {
    GetElementRequest request;
    GetElementResult result;
    request.set_task_id(task_id);
    TF_RETURN_IF_ERROR(client.GetElement(request, result));
    return result;
  }

  std::string GetDispatcherAddress() const {
    return test_cluster_->DispatcherAddress();
  }

  std::string GetWorkerAddress() const {
    return test_cluster_->WorkerAddress(0);
  }

  std::unique_ptr<TestCluster> test_cluster_;
  std::unique_ptr<DataServiceDispatcherClient> dispatcher_client_;
};

TEST_F(WorkerClientTest, LocalRead) {
  const int64 range = 5;
  TF_ASSERT_OK_AND_ASSIGN(const int64 dataset_id, RegisterDataset(range));
  TF_ASSERT_OK_AND_ASSIGN(const int64 job_client_id, CreateJob(dataset_id));
  TF_ASSERT_OK_AND_ASSIGN(const int64 task_id, GetTaskToRead(job_client_id));
  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<DataServiceWorkerClient> client,
                          GetWorkerClient(kLocalTransferProtocol));
  for (int64 i = 0; i < range; ++i) {
    TF_ASSERT_OK_AND_ASSIGN(GetElementResult result,
                            GetElement(*client, task_id));
    test::ExpectEqual(result.components[0], Tensor(int64{i * i}));
    EXPECT_FALSE(result.end_of_sequence);
  }

  // Remove the local worker from `LocalWorkers`. Since the client reads from a
  // local server, this should cause the request to fail.
  LocalWorkers::Remove(GetWorkerAddress());
  EXPECT_THAT(
      GetElement(*client, task_id),
      testing::StatusIs(error::CANCELLED,
                        MatchesRegex("Worker.*is no longer available.*")));
}

TEST_F(WorkerClientTest, LocalReadEmptyDataset) {
  TF_ASSERT_OK_AND_ASSIGN(const int64 dataset_id, RegisterDataset(/*range=*/0));
  TF_ASSERT_OK_AND_ASSIGN(const int64 job_client_id, CreateJob(dataset_id));
  TF_ASSERT_OK_AND_ASSIGN(const int64 task_id, GetTaskToRead(job_client_id));
  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<DataServiceWorkerClient> client,
                          GetWorkerClient(kLocalTransferProtocol));
  TF_ASSERT_OK_AND_ASSIGN(GetElementResult result,
                          GetElement(*client, task_id));
  EXPECT_TRUE(result.end_of_sequence);

  // Remove the local worker from `LocalWorkers`. Since the client reads from a
  // local server, this should cause the request to fail.
  LocalWorkers::Remove(GetWorkerAddress());
  EXPECT_THAT(
      GetElement(*client, task_id),
      testing::StatusIs(error::CANCELLED,
                        MatchesRegex("Worker.*is no longer available.*")));
}

TEST_F(WorkerClientTest, GrpcRead) {
  const int64 range = 5;
  TF_ASSERT_OK_AND_ASSIGN(const int64 dataset_id, RegisterDataset(range));
  TF_ASSERT_OK_AND_ASSIGN(const int64 job_client_id, CreateJob(dataset_id));
  TF_ASSERT_OK_AND_ASSIGN(const int64 task_id, GetTaskToRead(job_client_id));
  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<DataServiceWorkerClient> client,
                          GetWorkerClient(kGrpcTransferProtocol));
  for (int64 i = 0; i < range; ++i) {
    TF_ASSERT_OK_AND_ASSIGN(GetElementResult result,
                            GetElement(*client, task_id));
    test::ExpectEqual(result.components[0], Tensor(int64{i * i}));
    EXPECT_FALSE(result.end_of_sequence);
  }

  // Remove the local worker from `LocalWorkers`. Since the client reads from
  // gRPC, this will not cause the request to fail.
  LocalWorkers::Remove(GetWorkerAddress());
  TF_ASSERT_OK_AND_ASSIGN(GetElementResult result,
                          GetElement(*client, task_id));
  EXPECT_TRUE(result.end_of_sequence);
}

TEST_F(WorkerClientTest, LocalServerShutsDown) {
  TF_ASSERT_OK_AND_ASSIGN(const int64 dataset_id, RegisterDataset(/*range=*/5));
  TF_ASSERT_OK_AND_ASSIGN(const int64 job_client_id, CreateJob(dataset_id));
  TF_ASSERT_OK_AND_ASSIGN(const int64 task_id, GetTaskToRead(job_client_id));
  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<DataServiceWorkerClient> client,
                          GetWorkerClient(kLocalTransferProtocol));

  // Stopping a worker causes local reads to return Cancelled status.
  test_cluster_->StopWorkers();
  EXPECT_THAT(
      GetElement(*client, task_id),
      testing::StatusIs(error::CANCELLED,
                        MatchesRegex("Worker.*is no longer available.*")));
}

TEST_F(WorkerClientTest, CancelClient) {
  TF_ASSERT_OK_AND_ASSIGN(const int64 dataset_id, RegisterDataset(/*range=*/5));
  TF_ASSERT_OK_AND_ASSIGN(const int64 job_client_id, CreateJob(dataset_id));
  TF_ASSERT_OK_AND_ASSIGN(const int64 task_id, GetTaskToRead(job_client_id));
  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<DataServiceWorkerClient> client,
                          GetWorkerClient(kLocalTransferProtocol));

  client->TryCancel();
  EXPECT_THAT(GetElement(*client, task_id),
              testing::StatusIs(
                  error::CANCELLED,
                  MatchesRegex("Client for worker.*has been cancelled.")));
}

}  // namespace
}  // namespace data
}  // namespace tensorflow
