geode-for-redis/src/distributedTest/java/org/apache/geode/redis/internal/commands/executor/sortedset/ZRemRangeByScoreDUnitTest.java - geode - Git at Google

 /*
  * 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.internal.commands.executor.sortedset;

 import static java.util.Objects.requireNonNull;
 import static org.apache.geode.test.awaitility.GeodeAwaitility.await;
 import static org.apache.geode.test.dunit.rules.RedisClusterStartupRule.BIND_ADDRESS;
 import static org.apache.geode.test.dunit.rules.RedisClusterStartupRule.REDIS_CLIENT_TIMEOUT;
 import static org.assertj.core.api.Assertions.assertThat;

 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Future;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.stream.Collectors;

 import org.junit.After;
 import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
 import redis.clients.jedis.HostAndPort;
 import redis.clients.jedis.JedisCluster;
 import redis.clients.jedis.exceptions.JedisClusterMaxAttemptsException;

 import org.apache.geode.cache.Operation;
 import org.apache.geode.internal.AvailablePortHelper;
 import org.apache.geode.internal.cache.PartitionedRegion;
 import org.apache.geode.internal.cache.PartitionedRegionHelper;
 import org.apache.geode.redis.ConcurrentLoopingThreads;
 import org.apache.geode.redis.internal.data.RedisKey;
 import org.apache.geode.redis.internal.services.RegionProvider;
 import org.apache.geode.test.dunit.rules.ClusterStartupRule;
 import org.apache.geode.test.dunit.rules.MemberVM;
 import org.apache.geode.test.dunit.rules.RedisClusterStartupRule;
 import org.apache.geode.test.junit.rules.ExecutorServiceRule;

 public class ZRemRangeByScoreDUnitTest {
   @Rule
   public ExecutorServiceRule executor = new ExecutorServiceRule();

   @Rule
   public RedisClusterStartupRule clusterStartUp = new RedisClusterStartupRule(4);

   private JedisCluster jedis;
   private List<MemberVM> servers;
   private static final String KEY = "ZRemRangeByScoreDUnitTestKey";
   private static final String BASE_MEMBER_NAME = "member";
   private final int SET_SIZE = 500;
   private final AtomicBoolean isCrashing = new AtomicBoolean(false);

   @Before
   public void setup() {
     MemberVM locator = clusterStartUp.startLocatorVM(0);
     int locatorPort = locator.getPort();
     int[] redisPorts = AvailablePortHelper.getRandomAvailableTCPPorts(3);
     MemberVM server1 = clusterStartUp.startRedisVM(1, redisPorts[0], locatorPort);
     MemberVM server2 = clusterStartUp.startRedisVM(2, redisPorts[1], locatorPort);
     MemberVM server3 = clusterStartUp.startRedisVM(3, redisPorts[2], locatorPort);
     servers = new ArrayList<>();
     servers.add(server1);
     servers.add(server2);
     servers.add(server3);

     int redisServerPort1 = clusterStartUp.getRedisPort(1);

     jedis =
         new JedisCluster(new HostAndPort(BIND_ADDRESS, redisServerPort1), REDIS_CLIENT_TIMEOUT);
     isCrashing.set(false);
   }

   @After
   public void tearDown() {
     jedis.close();
   }

   @Test
   public void zRemRangeByScore_removesMembersInRangeFromBothServers() {
     Map<String, Double> memberScoreMap = makeMemberScoreMap();
     jedis.zadd(KEY, memberScoreMap);
     verifyDataExists(memberScoreMap);

     int removeRangeSize = 4;
     List<String> expected = new ArrayList<>(memberScoreMap.keySet().stream().sorted()
         .collect(Collectors.toList())).subList(removeRangeSize, SET_SIZE);

     // Subtract 1 from removeRangeSize here since scores start from 0, not 1
     assertThat(jedis.zremrangeByScore(KEY, "0", String.valueOf(removeRangeSize - 1)))
         .isEqualTo(removeRangeSize);

     for (int i = 0; i < removeRangeSize; ++i) {
       assertThat(jedis.zrank(KEY, BASE_MEMBER_NAME + i)).isNull();
     }

     assertThat(jedis.zrange(KEY, 0, -1)).containsExactlyElementsOf(expected);
   }

   @Test
   public void zRemRangeByScore_concurrentlyRemovesMembersInRangeFromBothServers() {
     Map<String, Double> memberScoreMap = makeMemberScoreMap();
     jedis.zadd(KEY, memberScoreMap);
     verifyDataExists(memberScoreMap);

     AtomicInteger totalRemoved = new AtomicInteger();
     int rangeSize = 5;
     new ConcurrentLoopingThreads(SET_SIZE / rangeSize,
         (i) -> doZRemRangeByScore(i * rangeSize, totalRemoved),
         (i) -> doZRemRangeByScore(i * rangeSize, totalRemoved)).run();

     assertThat(jedis.exists(KEY)).isFalse();

     assertThat(totalRemoved.get()).isEqualTo(SET_SIZE);
   }

   @Test
   public void zRemRangeByScoreRemovesMembersFromSortedSetAfterPrimaryShutsDown() {
     Map<String, Double> memberScoreMap = makeMemberScoreMap();
     jedis.zadd(KEY, memberScoreMap);
     verifyDataExists(memberScoreMap);

     stopNodeWithPrimaryBucketOfTheKey(false);

     doZRemRangeByScoreWithRetries(memberScoreMap);

     verifyDataDoesNotExist(memberScoreMap);
     assertThat(jedis.exists(KEY)).isFalse();
   }

   @Test
   public void zRemRangeByScoreCanRemoveMembersFromSortedSetWhenPrimaryIsCrashed()
       throws ExecutionException, InterruptedException {
     Map<String, Double> memberScoreMap = makeMemberScoreMap();

     jedis.zadd(KEY, memberScoreMap);
     verifyDataExists(memberScoreMap);

     String firstMember = String.join("", jedis.zrange(KEY, 0, 0));
     memberScoreMap.remove(firstMember);

     Future<Void> future1 = executor.submit(() -> stopNodeWithPrimaryBucketOfTheKey(true));
     Future<Void> future2 = executor.submit(this::removeAllButFirstEntry);

     future1.get();
     future2.get();

     await().until(() -> verifyDataDoesNotExist(memberScoreMap));
     assertThat(jedis.zrank(KEY, firstMember)).isZero();
     assertThat(jedis.exists(KEY)).isTrue();
   }

   private void verifyDataExists(Map<String, Double> memberScoreMap) {
     for (String member : memberScoreMap.keySet()) {
       Double score = jedis.zscore(KEY, member);
       assertThat(score).isEqualTo(memberScoreMap.get(member));
     }
   }

   private void doZRemRangeByScore(int i, AtomicInteger total) {
     long count = jedis.zremrangeByScore(KEY, String.valueOf(i), String.valueOf(i + 5));
     total.addAndGet((int) count);
   }

   private void stopNodeWithPrimaryBucketOfTheKey(boolean isCrash) {
     boolean isPrimary;
     for (MemberVM server : servers) {
       isPrimary = server.invoke(ZRemRangeByScoreDUnitTest::isPrimaryForKey);
       if (isPrimary) {
         if (isCrash) {
           isCrashing.set(true);
           server.getVM().bounceForcibly();
         } else {
           server.stop();
         }
         return;
       }
     }
   }

   private boolean verifyDataDoesNotExist(Map<String, Double> memberScoreMap) {
     try {
       for (String member : memberScoreMap.keySet()) {
         Double score = jedis.zscore(KEY, member);
         assertThat(score).isNull();
       }
     } catch (JedisClusterMaxAttemptsException e) {
       return false;
     }
     return true;
   }

   private void removeAllButFirstEntry() {
     long removed = 0;
     int rangeSize = 10;
     await().until(isCrashing::get);
     for (int i = 1; i < SET_SIZE; i += rangeSize) {
       removed += jedis.zremrangeByScore(KEY, String.valueOf(i), String.valueOf(i + rangeSize));
     }
     assertThat(removed).isEqualTo(SET_SIZE - 1);
   }

   private static boolean isPrimaryForKey() {
     PartitionedRegion region = (PartitionedRegion) requireNonNull(ClusterStartupRule.getCache())
         .getRegion(RegionProvider.DEFAULT_REDIS_REGION_NAME);
     int bucketId = PartitionedRegionHelper
         .getHashKey(region, Operation.GET, new RedisKey(KEY.getBytes()), null, null);
     return region.getLocalPrimaryBucketsListTestOnly().contains(bucketId);
   }

   private Map<String, Double> makeMemberScoreMap() {
     Map<String, Double> scoreMemberPairs = new HashMap<>();
     String memberName = BASE_MEMBER_NAME;
     for (int i = 0; i < SET_SIZE; i++) {
       memberName = memberName + i;
       scoreMemberPairs.put(memberName, (double) i);
     }
     return scoreMemberPairs;
   }

   private void doZRemRangeByScoreWithRetries(Map<String, Double> map) {
     int maxRetryAttempts = 10;
     int retryAttempts = 0;
     while (!zRemRangeByScoreWithRetries(map, retryAttempts, maxRetryAttempts)) {
       retryAttempts++;
     }
   }

   private boolean zRemRangeByScoreWithRetries(Map<String, Double> map, int retries,
       int maxRetries) {
     long removed;
     try {
       removed = jedis.zremrangeByScore(KEY, "-inf", "+inf");
     } catch (JedisClusterMaxAttemptsException e) {
       if (retries < maxRetries) {
         return false;
       }
       throw e;
     }
     assertThat(removed).as("Map for key " + KEY + ", removed is " + removed + " but map is size "
         + map.size() + ", retries is " + retries).isEqualTo(map.size());

     return true;
   }
 }
	/*
	* 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.internal.commands.executor.sortedset;

	import static java.util.Objects.requireNonNull;
	import static org.apache.geode.test.awaitility.GeodeAwaitility.await;
	import static org.apache.geode.test.dunit.rules.RedisClusterStartupRule.BIND_ADDRESS;
	import static org.apache.geode.test.dunit.rules.RedisClusterStartupRule.REDIS_CLIENT_TIMEOUT;
	import static org.assertj.core.api.Assertions.assertThat;

	import java.util.ArrayList;
	import java.util.HashMap;
	import java.util.List;
	import java.util.Map;
	import java.util.concurrent.ExecutionException;
	import java.util.concurrent.Future;
	import java.util.concurrent.atomic.AtomicBoolean;
	import java.util.concurrent.atomic.AtomicInteger;
	import java.util.stream.Collectors;

	import org.junit.After;
	import org.junit.Before;
	import org.junit.Rule;
	import org.junit.Test;
	import redis.clients.jedis.HostAndPort;
	import redis.clients.jedis.JedisCluster;
	import redis.clients.jedis.exceptions.JedisClusterMaxAttemptsException;

	import org.apache.geode.cache.Operation;
	import org.apache.geode.internal.AvailablePortHelper;
	import org.apache.geode.internal.cache.PartitionedRegion;
	import org.apache.geode.internal.cache.PartitionedRegionHelper;
	import org.apache.geode.redis.ConcurrentLoopingThreads;
	import org.apache.geode.redis.internal.data.RedisKey;
	import org.apache.geode.redis.internal.services.RegionProvider;
	import org.apache.geode.test.dunit.rules.ClusterStartupRule;
	import org.apache.geode.test.dunit.rules.MemberVM;
	import org.apache.geode.test.dunit.rules.RedisClusterStartupRule;
	import org.apache.geode.test.junit.rules.ExecutorServiceRule;

	public class ZRemRangeByScoreDUnitTest {
	@Rule
	public ExecutorServiceRule executor = new ExecutorServiceRule();

	@Rule
	public RedisClusterStartupRule clusterStartUp = new RedisClusterStartupRule(4);

	private JedisCluster jedis;
	private List<MemberVM> servers;
	private static final String KEY = "ZRemRangeByScoreDUnitTestKey";
	private static final String BASE_MEMBER_NAME = "member";
	private final int SET_SIZE = 500;
	private final AtomicBoolean isCrashing = new AtomicBoolean(false);

	@Before
	public void setup() {
	MemberVM locator = clusterStartUp.startLocatorVM(0);
	int locatorPort = locator.getPort();
	int[] redisPorts = AvailablePortHelper.getRandomAvailableTCPPorts(3);
	MemberVM server1 = clusterStartUp.startRedisVM(1, redisPorts[0], locatorPort);
	MemberVM server2 = clusterStartUp.startRedisVM(2, redisPorts[1], locatorPort);
	MemberVM server3 = clusterStartUp.startRedisVM(3, redisPorts[2], locatorPort);
	servers = new ArrayList<>();
	servers.add(server1);
	servers.add(server2);
	servers.add(server3);

	int redisServerPort1 = clusterStartUp.getRedisPort(1);

	jedis =
	new JedisCluster(new HostAndPort(BIND_ADDRESS, redisServerPort1), REDIS_CLIENT_TIMEOUT);
	isCrashing.set(false);
	}

	@After
	public void tearDown() {
	jedis.close();
	}

	@Test
	public void zRemRangeByScore_removesMembersInRangeFromBothServers() {
	Map<String, Double> memberScoreMap = makeMemberScoreMap();
	jedis.zadd(KEY, memberScoreMap);
	verifyDataExists(memberScoreMap);

	int removeRangeSize = 4;
	List<String> expected = new ArrayList<>(memberScoreMap.keySet().stream().sorted()
	.collect(Collectors.toList())).subList(removeRangeSize, SET_SIZE);

	// Subtract 1 from removeRangeSize here since scores start from 0, not 1
	assertThat(jedis.zremrangeByScore(KEY, "0", String.valueOf(removeRangeSize - 1)))
	.isEqualTo(removeRangeSize);

	for (int i = 0; i < removeRangeSize; ++i) {
	assertThat(jedis.zrank(KEY, BASE_MEMBER_NAME + i)).isNull();
	}

	assertThat(jedis.zrange(KEY, 0, -1)).containsExactlyElementsOf(expected);
	}

	@Test
	public void zRemRangeByScore_concurrentlyRemovesMembersInRangeFromBothServers() {
	Map<String, Double> memberScoreMap = makeMemberScoreMap();
	jedis.zadd(KEY, memberScoreMap);
	verifyDataExists(memberScoreMap);

	AtomicInteger totalRemoved = new AtomicInteger();
	int rangeSize = 5;
	new ConcurrentLoopingThreads(SET_SIZE / rangeSize,
	(i) -> doZRemRangeByScore(i * rangeSize, totalRemoved),
	(i) -> doZRemRangeByScore(i * rangeSize, totalRemoved)).run();

	assertThat(jedis.exists(KEY)).isFalse();

	assertThat(totalRemoved.get()).isEqualTo(SET_SIZE);
	}

	@Test
	public void zRemRangeByScoreRemovesMembersFromSortedSetAfterPrimaryShutsDown() {
	Map<String, Double> memberScoreMap = makeMemberScoreMap();
	jedis.zadd(KEY, memberScoreMap);
	verifyDataExists(memberScoreMap);

	stopNodeWithPrimaryBucketOfTheKey(false);

	doZRemRangeByScoreWithRetries(memberScoreMap);

	verifyDataDoesNotExist(memberScoreMap);
	assertThat(jedis.exists(KEY)).isFalse();
	}

	@Test
	public void zRemRangeByScoreCanRemoveMembersFromSortedSetWhenPrimaryIsCrashed()
	throws ExecutionException, InterruptedException {
	Map<String, Double> memberScoreMap = makeMemberScoreMap();

	jedis.zadd(KEY, memberScoreMap);
	verifyDataExists(memberScoreMap);

	String firstMember = String.join("", jedis.zrange(KEY, 0, 0));
	memberScoreMap.remove(firstMember);

	Future<Void> future1 = executor.submit(() -> stopNodeWithPrimaryBucketOfTheKey(true));
	Future<Void> future2 = executor.submit(this::removeAllButFirstEntry);

	future1.get();
	future2.get();

	await().until(() -> verifyDataDoesNotExist(memberScoreMap));
	assertThat(jedis.zrank(KEY, firstMember)).isZero();
	assertThat(jedis.exists(KEY)).isTrue();
	}

	private void verifyDataExists(Map<String, Double> memberScoreMap) {
	for (String member : memberScoreMap.keySet()) {
	Double score = jedis.zscore(KEY, member);
	assertThat(score).isEqualTo(memberScoreMap.get(member));
	}
	}

	private void doZRemRangeByScore(int i, AtomicInteger total) {
	long count = jedis.zremrangeByScore(KEY, String.valueOf(i), String.valueOf(i + 5));
	total.addAndGet((int) count);
	}

	private void stopNodeWithPrimaryBucketOfTheKey(boolean isCrash) {
	boolean isPrimary;
	for (MemberVM server : servers) {
	isPrimary = server.invoke(ZRemRangeByScoreDUnitTest::isPrimaryForKey);
	if (isPrimary) {
	if (isCrash) {
	isCrashing.set(true);
	server.getVM().bounceForcibly();
	} else {
	server.stop();
	}
	return;
	}
	}
	}

	private boolean verifyDataDoesNotExist(Map<String, Double> memberScoreMap) {
	try {
	for (String member : memberScoreMap.keySet()) {
	Double score = jedis.zscore(KEY, member);
	assertThat(score).isNull();
	}
	} catch (JedisClusterMaxAttemptsException e) {
	return false;
	}
	return true;
	}

	private void removeAllButFirstEntry() {
	long removed = 0;
	int rangeSize = 10;
	await().until(isCrashing::get);
	for (int i = 1; i < SET_SIZE; i += rangeSize) {
	removed += jedis.zremrangeByScore(KEY, String.valueOf(i), String.valueOf(i + rangeSize));
	}
	assertThat(removed).isEqualTo(SET_SIZE - 1);
	}

	private static boolean isPrimaryForKey() {
	PartitionedRegion region = (PartitionedRegion) requireNonNull(ClusterStartupRule.getCache())
	.getRegion(RegionProvider.DEFAULT_REDIS_REGION_NAME);
	int bucketId = PartitionedRegionHelper
	.getHashKey(region, Operation.GET, new RedisKey(KEY.getBytes()), null, null);
	return region.getLocalPrimaryBucketsListTestOnly().contains(bucketId);
	}

	private Map<String, Double> makeMemberScoreMap() {
	Map<String, Double> scoreMemberPairs = new HashMap<>();
	String memberName = BASE_MEMBER_NAME;
	for (int i = 0; i < SET_SIZE; i++) {
	memberName = memberName + i;
	scoreMemberPairs.put(memberName, (double) i);
	}
	return scoreMemberPairs;
	}

	private void doZRemRangeByScoreWithRetries(Map<String, Double> map) {
	int maxRetryAttempts = 10;
	int retryAttempts = 0;
	while (!zRemRangeByScoreWithRetries(map, retryAttempts, maxRetryAttempts)) {
	retryAttempts++;
	}
	}

	private boolean zRemRangeByScoreWithRetries(Map<String, Double> map, int retries,
	int maxRetries) {
	long removed;
	try {
	removed = jedis.zremrangeByScore(KEY, "-inf", "+inf");
	} catch (JedisClusterMaxAttemptsException e) {
	if (retries < maxRetries) {
	return false;
	}
	throw e;
	}
	assertThat(removed).as("Map for key " + KEY + ", removed is " + removed + " but map is size "
	+ map.size() + ", retries is " + retries).isEqualTo(map.size());

	return true;
	}
	}