/*
 * The Alluxio Open Foundation licenses this work under the Apache License, version 2.0
 * (the "License"). You may not use this work except in compliance with the License, which is
 * available at www.apache.org/licenses/LICENSE-2.0
 *
 * This software is distributed on an "AS IS" basis, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND,
 * either express or implied, as more fully set forth in the License.
 *
 * See the NOTICE file distributed with this work for information regarding copyright ownership.
 */

package alluxio.master.block;

import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertTrue;

import alluxio.conf.Configuration;
import alluxio.conf.PropertyKey;
import alluxio.grpc.GetRegisterLeasePRequest;
import alluxio.util.SleepUtils;
import alluxio.wire.RegisterLease;

import org.junit.Before;
import org.junit.Test;

import java.util.Optional;

public class RegisterLeaseManagerTest {
  public static final long WORKER1_ID = 1L;
  public static final long WORKER2_ID = 2L;
  public static final long WORKER3_ID = 3L;
  public static final long WORKER4_ID = 4L;

  private RegisterLeaseManager mLeaseManager;

  @Before
  public void before() {
    Configuration.set(PropertyKey.MASTER_WORKER_REGISTER_LEASE_COUNT, 2);
    Configuration.set(PropertyKey.MASTER_WORKER_REGISTER_LEASE_TTL, "3s");
    // Tests on the JVM check logic will be done separately
    Configuration.set(PropertyKey.MASTER_WORKER_REGISTER_LEASE_RESPECT_JVM_SPACE, false);

    mLeaseManager = new RegisterLeaseManager();
  }

  @Test
  public void acquireVerifyRelease() {
    GetRegisterLeasePRequest request = GetRegisterLeasePRequest.newBuilder()
        .setWorkerId(WORKER1_ID).setBlockCount(0).build();

    Optional<RegisterLease> lease = mLeaseManager.tryAcquireLease(request);
    assertTrue(lease.isPresent());
    assertTrue(mLeaseManager.hasLease(WORKER1_ID));

    mLeaseManager.releaseLease(WORKER1_ID);
    assertFalse(mLeaseManager.hasLease(WORKER1_ID));
  }

  @Test
  public void recycleExpiredLease() {
    // allocate two leases
    GetRegisterLeasePRequest request1 = GetRegisterLeasePRequest.newBuilder()
        .setWorkerId(WORKER1_ID).setBlockCount(0).build();
    Optional<RegisterLease> lease1 = mLeaseManager.tryAcquireLease(request1);
    GetRegisterLeasePRequest request2 = GetRegisterLeasePRequest.newBuilder()
        .setWorkerId(WORKER2_ID).setBlockCount(0).build();
    Optional<RegisterLease> lease2 = mLeaseManager.tryAcquireLease(request2);
    assertTrue(lease1.isPresent());
    assertTrue(lease2.isPresent());
    assertTrue(mLeaseManager.hasLease(WORKER1_ID));
    assertTrue(mLeaseManager.hasLease(WORKER2_ID));

    // wait for expiration
    SleepUtils.sleepMs(5000);

    // allocate two more leases
    GetRegisterLeasePRequest request3 = GetRegisterLeasePRequest.newBuilder()
        .setWorkerId(WORKER3_ID).setBlockCount(500_000).build();
    Optional<RegisterLease> lease3 = mLeaseManager.tryAcquireLease(request3);
    GetRegisterLeasePRequest request4 = GetRegisterLeasePRequest.newBuilder()
        .setWorkerId(WORKER4_ID).setBlockCount(100_000).build();
    Optional<RegisterLease> lease4 = mLeaseManager.tryAcquireLease(request4);

    // The 2 old leases should be recycled
    assertFalse(mLeaseManager.hasLease(WORKER1_ID));
    assertFalse(mLeaseManager.hasLease(WORKER2_ID));
    assertTrue(lease3.isPresent());
    assertTrue(lease4.isPresent());
    assertTrue(mLeaseManager.hasLease(WORKER3_ID));
    assertTrue(mLeaseManager.hasLease(WORKER4_ID));

    // Acquiring one more lease before existing ones expire will be blocked
    GetRegisterLeasePRequest shouldWait = GetRegisterLeasePRequest.newBuilder()
        .setWorkerId(5L).setBlockCount(100_000).build();
    Optional<RegisterLease> empty = mLeaseManager.tryAcquireLease(shouldWait);
    assertFalse(empty.isPresent());
    assertFalse(mLeaseManager.hasLease(5L));
  }

  @Test
  public void findExistingLease() {
    GetRegisterLeasePRequest request1 = GetRegisterLeasePRequest.newBuilder()
        .setWorkerId(WORKER1_ID).setBlockCount(0).build();
    Optional<RegisterLease> lease1 = mLeaseManager.tryAcquireLease(request1);
    assertTrue(lease1.isPresent());
    assertTrue(mLeaseManager.hasLease(WORKER1_ID));

    // Same worker requesting for an existing lease
    Optional<RegisterLease> sameLease = mLeaseManager.tryAcquireLease(request1);
    assertTrue(sameLease.isPresent());
    assertTrue(mLeaseManager.hasLease(WORKER1_ID));

    // The same lease is granted so there should be another vacancy
    GetRegisterLeasePRequest request2 = GetRegisterLeasePRequest.newBuilder()
        .setWorkerId(WORKER2_ID).setBlockCount(0).build();
    Optional<RegisterLease> lease2 = mLeaseManager.tryAcquireLease(request2);
    assertTrue(lease2.isPresent());
    assertTrue(mLeaseManager.hasLease(WORKER2_ID));
    // The lease for worker 1 is not recycled
    assertTrue(mLeaseManager.hasLease(WORKER1_ID));

    // Now all token have been given out, new ones will have to wait
    GetRegisterLeasePRequest shouldWait = GetRegisterLeasePRequest.newBuilder()
        .setWorkerId(5L).setBlockCount(100_000).build();
    Optional<RegisterLease> empty = mLeaseManager.tryAcquireLease(shouldWait);
    assertFalse(empty.isPresent());
    assertFalse(mLeaseManager.hasLease(5L));
  }
}
