geode-redis/src/distributedTest/java/org/apache/geode/redis/CheckPrimaryBucketFunction.java

/*
 * 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.geode.redis;

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

import org.apache.logging.log4j.Logger;

import org.apache.geode.cache.Region;
import org.apache.geode.cache.execute.Function;
import org.apache.geode.cache.execute.FunctionContext;
import org.apache.geode.cache.execute.ResultSender;
import org.apache.geode.cache.partition.PartitionRegionHelper;
import org.apache.geode.distributed.DistributedMember;
import org.apache.geode.internal.cache.LocalDataSet;
import org.apache.geode.internal.cache.PartitionedRegion;
import org.apache.geode.internal.cache.PrimaryBucketLockException;
import org.apache.geode.internal.cache.execute.BucketMovedException;
import org.apache.geode.internal.cache.execute.RegionFunctionContextImpl;
import org.apache.geode.logging.internal.log4j.api.LogService;
import org.apache.geode.test.awaitility.GeodeAwaitility;

@SuppressWarnings("unchecked")
public class CheckPrimaryBucketFunction implements Function {
  public static final String ID = CheckPrimaryBucketFunction.class.getName();
  private final CountDownLatch signalFunctionHasStarted = new CountDownLatch(1);
  private final CountDownLatch signalPrimaryHasMoved = new CountDownLatch(1);
  private static final Logger logger = LogService.getLogger();

  public void waitForFunctionToStart() {
    try {
      signalFunctionHasStarted.await();
    } catch (InterruptedException e) {
      e.printStackTrace();
    }
  }

  public void finishedMovingPrimary() {
    signalPrimaryHasMoved.countDown();
  }

  @Override
  public void execute(FunctionContext context) {
    RegionFunctionContextImpl regionFunctionContext = (RegionFunctionContextImpl) context;
    String key = (String) regionFunctionContext.getFilter().iterator().next();
    boolean releaseLatchEarly = (boolean) context.getArguments();

    ResultSender result = context.getResultSender();
    DistributedMember member = context.getCache().getDistributedSystem().getDistributedMember();

    if (!isMemberPrimary(regionFunctionContext, key, member)) {
      LogService.getLogger().error("Member is not primary.");
      result.lastResult(false);
      return;
    }

    Region<?, ?> localRegion =
        regionFunctionContext.getLocalDataSet(regionFunctionContext.getDataSet());

    if (releaseLatchEarly) {
      signalFunctionHasStarted.countDown();
      // now wait until test has moved primary
      try {
        signalPrimaryHasMoved.await();
      } catch (InterruptedException e) {
        throw new IllegalStateException(e);
      }
    }

    Runnable r = () -> {
      if (!releaseLatchEarly) {
        signalFunctionHasStarted.countDown();
      }

      try {
        GeodeAwaitility.await()
            .during(10, TimeUnit.SECONDS)
            .atMost(11, TimeUnit.SECONDS)
            .until(() -> isMemberPrimary(regionFunctionContext, key, member));
        result.lastResult(true);
      } catch (Exception e) {
        e.printStackTrace();
        result.lastResult(false);
      }
    };

    LocalDataSet localDataSet = (LocalDataSet) localRegion;
    PartitionedRegion partitionedRegion = localDataSet.getProxy();
    try {
      partitionedRegion.computeWithPrimaryLocked(key, r);
    } catch (PrimaryBucketLockException ex) {
      throw new BucketMovedException(ex.toString());
    }
  }

  private boolean isMemberPrimary(RegionFunctionContextImpl context, String key,
      DistributedMember member) {
    DistributedMember primaryForKey = PartitionRegionHelper
        .getPrimaryMemberForKey(context.getDataSet(), key);

    return primaryForKey.equals(member);
  }

  @Override
  public boolean optimizeForWrite() {
    return true;
  }

  @Override
  public boolean isHA() {
    return true;
  }

  @Override
  public String getId() {
    return ID;
  }
}