/**
 * 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.
 */

package org.apache.tajo.master.rm;

import com.google.protobuf.RpcCallback;
import org.apache.hadoop.yarn.proto.YarnProtos;
import org.apache.tajo.QueryId;
import org.apache.tajo.QueryIdFactory;
import org.apache.tajo.conf.TajoConf;
import org.apache.tajo.ipc.TajoMasterProtocol.*;
import org.apache.tajo.rpc.NullCallback;
import org.apache.tajo.rpc.protocolrecords.PrimitiveProtos;
import org.junit.Test;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;

import static org.apache.tajo.ipc.TajoResourceTrackerProtocol.NodeHeartbeat;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertTrue;

public class TestTajoResourceManager {
  private final PrimitiveProtos.BoolProto BOOL_TRUE = PrimitiveProtos.BoolProto.newBuilder().setValue(true).build();
  private final PrimitiveProtos.BoolProto BOOL_FALSE = PrimitiveProtos.BoolProto.newBuilder().setValue(false).build();

  TajoConf tajoConf;

  long queryIdTime = System.currentTimeMillis();
  int numWorkers = 5;
  float workerDiskSlots = 5.0f;
  int workerMemoryMB = 512 * 10;
  WorkerResourceAllocationResponse response;

  private TajoWorkerResourceManager initResourceManager(boolean queryMasterMode) throws Exception {
    tajoConf = new org.apache.tajo.conf.TajoConf();

    tajoConf.setFloatVar(TajoConf.ConfVars.TAJO_QUERYMASTER_DISK_SLOT, 0.0f);
    tajoConf.setIntVar(TajoConf.ConfVars.TAJO_QUERYMASTER_MEMORY_MB, 512);
    tajoConf.setVar(TajoConf.ConfVars.RESOURCE_TRACKER_RPC_ADDRESS, "localhost:0");
    TajoWorkerResourceManager tajoWorkerResourceManager = new TajoWorkerResourceManager(tajoConf);
    tajoWorkerResourceManager.init(tajoConf);
    tajoWorkerResourceManager.start();

    for(int i = 0; i < numWorkers; i++) {
      ServerStatusProto.System system = ServerStatusProto.System.newBuilder()
          .setAvailableProcessors(1)
          .setFreeMemoryMB(workerMemoryMB)
          .setMaxMemoryMB(workerMemoryMB)
          .setTotalMemoryMB(workerMemoryMB)
          .build();

      ServerStatusProto.JvmHeap jvmHeap = ServerStatusProto.JvmHeap.newBuilder()
          .setFreeHeap(workerMemoryMB)
          .setMaxHeap(workerMemoryMB)
          .setTotalHeap(workerMemoryMB)
          .build();

      ServerStatusProto.Disk disk = ServerStatusProto.Disk.newBuilder()
          .setAbsolutePath("/")
          .setFreeSpace(0)
          .setTotalSpace(0)
          .setUsableSpace(0)
          .build();

      List<ServerStatusProto.Disk> disks = new ArrayList<ServerStatusProto.Disk>();

      disks.add(disk);

      ServerStatusProto serverStatus = ServerStatusProto.newBuilder()
          .setQueryMasterMode(queryMasterMode ? BOOL_TRUE : BOOL_FALSE)
          .setTaskRunnerMode(BOOL_TRUE)
          .setDiskSlots(workerDiskSlots)
          .setMemoryResourceMB(workerMemoryMB)
          .setJvmHeap(jvmHeap)
          .setSystem(system)
          .addAllDisk(disks)
          .setRunningTaskNum(0)
          .build();

      NodeHeartbeat tajoHeartbeat = NodeHeartbeat.newBuilder()
          .setTajoWorkerHost("host" + (i + 1))
          .setTajoQueryMasterPort(21000)
          .setTajoWorkerHttpPort(28080 + i)
          .setPeerRpcPort(12345)
          .setServerStatus(serverStatus)
          .build();

      tajoWorkerResourceManager.getResourceTracker().heartbeat(null, tajoHeartbeat, NullCallback.get());
    }

    return tajoWorkerResourceManager;
  }


  @Test
  public void testHeartbeat() throws Exception {
    TajoWorkerResourceManager tajoWorkerResourceManager = null;
    try {
      tajoWorkerResourceManager = initResourceManager(false);
      assertEquals(numWorkers, tajoWorkerResourceManager.getWorkers().size());
      for(Worker worker: tajoWorkerResourceManager.getWorkers().values()) {
        WorkerResource resource = worker.getResource();
        assertEquals(workerMemoryMB, resource.getAvailableMemoryMB());
        assertEquals(workerDiskSlots, resource.getAvailableDiskSlots(), 0);
      }
    } finally {
      if (tajoWorkerResourceManager != null) {
        tajoWorkerResourceManager.stop();
      }
    }
  }

  @Test
  public void testMemoryResource() throws Exception {
    TajoWorkerResourceManager tajoWorkerResourceManager = null;
    try {
      tajoWorkerResourceManager = initResourceManager(false);

      final int minMemory = 256;
      final int maxMemory = 512;
      float diskSlots = 1.0f;

      QueryId queryId = QueryIdFactory.newQueryId(queryIdTime, 1);

      WorkerResourceAllocationRequest request = WorkerResourceAllocationRequest.newBuilder()
          .setResourceRequestPriority(ResourceRequestPriority.MEMORY)
          .setNumContainers(60)
          .setQueryId(queryId.getProto())
          .setMaxDiskSlotPerContainer(diskSlots)
          .setMinDiskSlotPerContainer(diskSlots)
          .setMinMemoryMBPerContainer(minMemory)
          .setMaxMemoryMBPerContainer(maxMemory)
          .build();

      final CountDownLatch barrier = new CountDownLatch(1);
      final List<YarnProtos.ContainerIdProto> containerIds = new ArrayList<YarnProtos.ContainerIdProto>();

      RpcCallback<WorkerResourceAllocationResponse> callBack = new RpcCallback<WorkerResourceAllocationResponse>() {

        @Override
        public void run(WorkerResourceAllocationResponse response) {
          TestTajoResourceManager.this.response = response;
          barrier.countDown();
        }
      };

      tajoWorkerResourceManager.allocateWorkerResources(request, callBack);
      assertTrue(barrier.await(3, TimeUnit.SECONDS));


      // assert after callback
      int totalUsedMemory = 0;
      int totalUsedDisks = 0;

      for(Worker worker: tajoWorkerResourceManager.getWorkers().values()) {
        WorkerResource resource = worker.getResource();
        assertEquals(0, resource.getAvailableMemoryMB());
        assertEquals(0, resource.getAvailableDiskSlots(), 0);
        assertEquals(5.0f, resource.getUsedDiskSlots(), 0);

        totalUsedMemory += resource.getUsedMemoryMB();
        totalUsedDisks += resource.getUsedDiskSlots();
      }

      assertEquals(workerMemoryMB * numWorkers, totalUsedMemory);
      assertEquals(workerDiskSlots * numWorkers, totalUsedDisks, 0);

      assertEquals(numWorkers * 10, response.getWorkerAllocatedResourceList().size());

      for(WorkerAllocatedResource eachResource: response.getWorkerAllocatedResourceList()) {
        assertTrue(
            eachResource.getAllocatedMemoryMB() >= minMemory &&  eachResource.getAllocatedMemoryMB() <= maxMemory);
        containerIds.add(eachResource.getContainerId());
      }

      for(YarnProtos.ContainerIdProto eachContainerId: containerIds) {
        tajoWorkerResourceManager.releaseWorkerResource(eachContainerId);
      }

      for(Worker worker: tajoWorkerResourceManager.getWorkers().values()) {
        WorkerResource resource = worker.getResource();
        assertEquals(workerMemoryMB, resource.getAvailableMemoryMB());
        assertEquals(0, resource.getUsedMemoryMB());

        assertEquals(workerDiskSlots, resource.getAvailableDiskSlots(), 0);
        assertEquals(0.0f, resource.getUsedDiskSlots(), 0);
      }
    } finally {
      if (tajoWorkerResourceManager != null) {
        tajoWorkerResourceManager.stop();
      }
    }
  }

  @Test
  public void testMemoryNotCommensurable() throws Exception {
    TajoWorkerResourceManager tajoWorkerResourceManager = null;

    try {
      tajoWorkerResourceManager = initResourceManager(false);

      final int minMemory = 200;
      final int maxMemory = 500;
      float diskSlots = 1.0f;

      QueryId queryId = QueryIdFactory.newQueryId(queryIdTime, 2);

      int requiredContainers = 60;

      int numAllocatedContainers = 0;

      int loopCount = 0;
      while(true) {
        WorkerResourceAllocationRequest request = WorkerResourceAllocationRequest.newBuilder()
            .setResourceRequestPriority(ResourceRequestPriority.MEMORY)
            .setNumContainers(requiredContainers - numAllocatedContainers)
            .setQueryId(queryId.getProto())
            .setMaxDiskSlotPerContainer(diskSlots)
            .setMinDiskSlotPerContainer(diskSlots)
            .setMinMemoryMBPerContainer(minMemory)
            .setMaxMemoryMBPerContainer(maxMemory)
            .build();

        final CountDownLatch barrier = new CountDownLatch(1);

        RpcCallback<WorkerResourceAllocationResponse> callBack = new RpcCallback<WorkerResourceAllocationResponse>() {
          @Override
          public void run(WorkerResourceAllocationResponse response) {
            TestTajoResourceManager.this.response = response;
            barrier.countDown();
          }
        };

        tajoWorkerResourceManager.allocateWorkerResources(request, callBack);

        assertTrue(barrier.await(3, TimeUnit.SECONDS));

        numAllocatedContainers += TestTajoResourceManager.this.response.getWorkerAllocatedResourceList().size();

        //release resource
        for(WorkerAllocatedResource eachResource:
            TestTajoResourceManager.this.response.getWorkerAllocatedResourceList()) {
          assertTrue(
              eachResource.getAllocatedMemoryMB() >= minMemory &&  eachResource.getAllocatedMemoryMB() <= maxMemory);
          tajoWorkerResourceManager.releaseWorkerResource(eachResource.getContainerId());
        }

        for(Worker worker: tajoWorkerResourceManager.getWorkers().values()) {
          WorkerResource resource = worker.getResource();
          assertEquals(0, resource.getUsedMemoryMB());
          assertEquals(workerMemoryMB, resource.getAvailableMemoryMB());

          assertEquals(0.0f, resource.getUsedDiskSlots(), 0);
          assertEquals(workerDiskSlots, resource.getAvailableDiskSlots(), 0);
        }

        loopCount++;

        if(loopCount == 2) {
          assertEquals(requiredContainers, numAllocatedContainers);
          break;
        }
      }

      for(Worker worker: tajoWorkerResourceManager.getWorkers().values()) {
        WorkerResource resource = worker.getResource();
        assertEquals(0, resource.getUsedMemoryMB());
        assertEquals(workerMemoryMB, resource.getAvailableMemoryMB());

        assertEquals(0.0f, resource.getUsedDiskSlots(), 0);
        assertEquals(workerDiskSlots, resource.getAvailableDiskSlots(), 0);
      }
    } finally {
      if (tajoWorkerResourceManager != null) {
        tajoWorkerResourceManager.stop();
      }
    }
  }

  @Test
  public void testDiskResource() throws Exception {
    TajoWorkerResourceManager tajoWorkerResourceManager = null;

    try {
      tajoWorkerResourceManager = initResourceManager(false);

      final float minDiskSlots = 1.0f;
      final float maxDiskSlots = 2.0f;
      int memoryMB = 256;

      QueryId queryId = QueryIdFactory.newQueryId(queryIdTime, 3);

      WorkerResourceAllocationRequest request = WorkerResourceAllocationRequest.newBuilder()
          .setResourceRequestPriority(ResourceRequestPriority.DISK)
          .setNumContainers(60)
          .setQueryId(queryId.getProto())
          .setMaxDiskSlotPerContainer(maxDiskSlots)
          .setMinDiskSlotPerContainer(minDiskSlots)
          .setMinMemoryMBPerContainer(memoryMB)
          .setMaxMemoryMBPerContainer(memoryMB)
          .build();

      final CountDownLatch barrier = new CountDownLatch(1);
      final List<YarnProtos.ContainerIdProto> containerIds = new ArrayList<YarnProtos.ContainerIdProto>();


      RpcCallback<WorkerResourceAllocationResponse> callBack = new RpcCallback<WorkerResourceAllocationResponse>() {

        @Override
        public void run(WorkerResourceAllocationResponse response) {
          TestTajoResourceManager.this.response = response;
          barrier.countDown();
        }
      };

      tajoWorkerResourceManager.allocateWorkerResources(request, callBack);
      assertTrue(barrier.await(3, TimeUnit.SECONDS));

      for(WorkerAllocatedResource eachResource: response.getWorkerAllocatedResourceList()) {
        assertTrue("AllocatedDiskSlot:" + eachResource.getAllocatedDiskSlots(),
            eachResource.getAllocatedDiskSlots() >= minDiskSlots &&
                eachResource.getAllocatedDiskSlots() <= maxDiskSlots);
        containerIds.add(eachResource.getContainerId());
      }

      // assert after callback
      int totalUsedDisks = 0;
      for(Worker worker: tajoWorkerResourceManager.getWorkers().values()) {
        WorkerResource resource = worker.getResource();
        //each worker allocated 3 container (2 disk slot = 2, 1 disk slot = 1)
        assertEquals(0, resource.getAvailableDiskSlots(), 0);
        assertEquals(5.0f, resource.getUsedDiskSlots(), 0);
        assertEquals(256 * 3, resource.getUsedMemoryMB());

        totalUsedDisks += resource.getUsedDiskSlots();
      }

      assertEquals(workerDiskSlots * numWorkers, totalUsedDisks, 0);

      assertEquals(numWorkers * 3, response.getWorkerAllocatedResourceList().size());

      for(YarnProtos.ContainerIdProto eachContainerId: containerIds) {
        tajoWorkerResourceManager.releaseWorkerResource(eachContainerId);
      }

      for(Worker worker: tajoWorkerResourceManager.getWorkers().values()) {
        WorkerResource resource = worker.getResource();
        assertEquals(workerMemoryMB, resource.getAvailableMemoryMB());
        assertEquals(0, resource.getUsedMemoryMB());

        assertEquals(workerDiskSlots, resource.getAvailableDiskSlots(), 0);
        assertEquals(0.0f, resource.getUsedDiskSlots(), 0);
      }
    } finally {
      if (tajoWorkerResourceManager != null) {
        tajoWorkerResourceManager.stop();
      }
    }
  }
}