[close #550] rawkv: fix seek leader/proxy store early abort (#551)

Co-authored-by: ti-srebot <66930949+ti-srebot@users.noreply.github.com>

src/main/java/org/tikv/common/region/AbstractRegionStoreClient.java

@@ -277,19 +277,27 @@ public abstract class AbstractRegionStoreClient
     List<SwitchLeaderTask> responses = new LinkedList<>();
     for (Metapb.Peer peer : region.getFollowerList()) {
       ByteString key = region.getStartKey();
-      TiStore peerStore = regionManager.getStoreById(peer.getStoreId(), backOffer);
+      try {
-      ManagedChannel channel =
+        TiStore peerStore = regionManager.getStoreById(peer.getStoreId(), backOffer);
-          channelFactory.getChannel(
+        ManagedChannel channel =
-              peerStore.getAddress(), regionManager.getPDClient().getHostMapping());
+            channelFactory.getChannel(
-      TikvGrpc.TikvFutureStub stub =
+                peerStore.getAddress(), regionManager.getPDClient().getHostMapping());
-          TikvGrpc.newFutureStub(channel).withDeadlineAfter(timeout, TimeUnit.MILLISECONDS);
+        TikvGrpc.TikvFutureStub stub =
-      Kvrpcpb.RawGetRequest rawGetRequest =
+            TikvGrpc.newFutureStub(channel).withDeadlineAfter(timeout, TimeUnit.MILLISECONDS);
-          Kvrpcpb.RawGetRequest.newBuilder()
+        Kvrpcpb.RawGetRequest rawGetRequest =
-              .setContext(region.getReplicaContext(peer))
+            Kvrpcpb.RawGetRequest.newBuilder()
-              .setKey(key)
+                .setContext(region.getReplicaContext(peer))
-              .build();
+                .setKey(key)
-      ListenableFuture<Kvrpcpb.RawGetResponse> task = stub.rawGet(rawGetRequest);
+                .build();
-      responses.add(new SwitchLeaderTask(task, peer));
+        ListenableFuture<Kvrpcpb.RawGetResponse> task = stub.rawGet(rawGetRequest);
+        responses.add(new SwitchLeaderTask(task, peer));
+      } catch (Exception e) {
+        logger.warn(
+            "switch region[{}] leader store to {} failed: {}",
+            region.getId(),
+            peer.getStoreId(),
+            e);
+      }
     }
     while (true) {
       try {
@@ -328,22 +336,31 @@ public abstract class AbstractRegionStoreClient
     List<ForwardCheckTask> responses = new LinkedList<>();
     for (Metapb.Peer peer : region.getFollowerList()) {
       ByteString key = region.getStartKey();
-      TiStore peerStore = regionManager.getStoreById(peer.getStoreId(), backOffer);
+      try {
-      ManagedChannel channel =
+        TiStore peerStore = regionManager.getStoreById(peer.getStoreId(), backOffer);
-          channelFactory.getChannel(
+        ManagedChannel channel =
-              peerStore.getAddress(), regionManager.getPDClient().getHostMapping());
+            channelFactory.getChannel(
-      TikvGrpc.TikvFutureStub stub =
+                peerStore.getAddress(), regionManager.getPDClient().getHostMapping());
-          TikvGrpc.newFutureStub(channel).withDeadlineAfter(forwardTimeout, TimeUnit.MILLISECONDS);
+        TikvGrpc.TikvFutureStub stub =
-      Metadata header = new Metadata();
+            TikvGrpc.newFutureStub(channel)
-      header.put(TiConfiguration.FORWARD_META_DATA_KEY, store.getStore().getAddress());
+                .withDeadlineAfter(forwardTimeout, TimeUnit.MILLISECONDS);
-      Kvrpcpb.RawGetRequest rawGetRequest =
+        Metadata header = new Metadata();
-          Kvrpcpb.RawGetRequest.newBuilder()
+        header.put(TiConfiguration.FORWARD_META_DATA_KEY, store.getStore().getAddress());
-              .setContext(region.getReplicaContext(region.getLeader()))
+        Kvrpcpb.RawGetRequest rawGetRequest =
-              .setKey(key)
+            Kvrpcpb.RawGetRequest.newBuilder()
-              .build();
+                .setContext(region.getReplicaContext(region.getLeader()))
-      ListenableFuture<Kvrpcpb.RawGetResponse> task =
+                .setKey(key)
-          MetadataUtils.attachHeaders(stub, header).rawGet(rawGetRequest);
+                .build();
-      responses.add(new ForwardCheckTask(task, peerStore.getStore()));
+        ListenableFuture<Kvrpcpb.RawGetResponse> task =
+            MetadataUtils.attachHeaders(stub, header).rawGet(rawGetRequest);
+        responses.add(new ForwardCheckTask(task, peerStore.getStore()));
+      } catch (Exception e) {
+        logger.warn(
+            "switch region[{}] leader store to {} failed: {}",
+            region.getId(),
+            peer.getStoreId(),
+            e);
+      }
     }
     while (true) {
       try {

src/main/java/org/tikv/common/region/RegionManager.java

@@ -26,7 +26,6 @@ import java.util.List;
 import java.util.concurrent.Executors;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.TimeUnit;
-import java.util.stream.Collectors;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.tikv.common.ReadOnlyPDClient;
@@ -47,6 +46,7 @@ import org.tikv.kvproto.Pdpb;
 
 @SuppressWarnings("UnstableApiUsage")
 public class RegionManager {
+
   private static final Logger logger = LoggerFactory.getLogger(RegionManager.class);
   public static final Histogram GET_REGION_BY_KEY_REQUEST_LATENCY =
       HistogramUtils.buildDuration()
@@ -205,22 +205,23 @@ public class RegionManager {
   }
 
   public TiRegion createRegion(Metapb.Region region, BackOffer backOffer) {
-    List<Metapb.Peer> peers = region.getPeersList();
+    return createRegion(region, null, backOffer);
-    List<TiStore> stores = getRegionStore(peers, backOffer);
-    return new TiRegion(conf, region, null, peers, stores);
   }
 
   private TiRegion createRegion(Metapb.Region region, Metapb.Peer leader, BackOffer backOffer) {
-    List<Metapb.Peer> peers = region.getPeersList();
+    List<Metapb.Peer> peers = new ArrayList<>();
-    List<TiStore> stores = getRegionStore(peers, backOffer);
+    List<TiStore> stores = new ArrayList<>();
-    return new TiRegion(conf, region, leader, peers, stores);
+    for (Metapb.Peer peer : region.getPeersList()) {
-  }
+      try {
-
+        stores.add(getStoreById(peer.getStoreId(), backOffer));
-  private List<TiStore> getRegionStore(List<Metapb.Peer> peers, BackOffer backOffer) {
+        peers.add(peer);
-    return peers
+      } catch (Exception e) {
-        .stream()
+        logger.warn("Store {} not found: {}", peer.getStoreId(), e.toString());
-        .map(p -> getStoreById(p.getStoreId(), backOffer))
+      }
-        .collect(Collectors.toList());
+    }
+    Metapb.Region newRegion =
+        Metapb.Region.newBuilder().mergeFrom(region).clearPeers().addAllPeers(peers).build();
+    return new TiRegion(conf, newRegion, leader, peers, stores);
   }
 
   private TiStore getStoreByIdWithBackOff(long id, BackOffer backOffer) {

src/test/java/org/tikv/common/PDClientMockTest.java

@@ -41,7 +41,7 @@ import org.tikv.kvproto.Metapb.StoreState;
 
 public class PDClientMockTest extends PDMockServerTest {
 
-  private static final String LOCAL_ADDR_IPV6 = "[::]";
+  private static final String LOCAL_ADDR_IPV6 = "[::1]";
   public static final String HTTP = "http://";
 
   @Test

src/test/java/org/tikv/common/PDMockServerTest.java

@@ -39,7 +39,7 @@ public abstract class PDMockServerTest {
 
   void setup(String addr) throws IOException {
     int basePort;
-    try (ServerSocket s = new ServerSocket(51820)) {
+    try (ServerSocket s = new ServerSocket(0)) {
       basePort = s.getLocalPort();
     }
 
@@ -54,9 +54,11 @@ public abstract class PDMockServerTest {
                   GrpcUtils.makeMember(2, "http://" + addr + ":" + (basePort + 1)),
                   GrpcUtils.makeMember(3, "http://" + addr + ":" + (basePort + 2))));
       pdServers.add(server);
+      if (i == 0) {
+        leader = server;
+      }
     }
 
-    leader = pdServers.get(0);
     TiConfiguration conf = TiConfiguration.createDefault(addr + ":" + leader.port);
     conf.setKvMode("RAW");
     conf.setWarmUpEnable(false);

src/test/java/org/tikv/common/SeekLeaderStoreTest.java

@@ -18,9 +18,13 @@
 package org.tikv.common;
 
 import com.google.protobuf.ByteString;
+import java.util.concurrent.atomic.AtomicInteger;
 import org.junit.Assert;
 import org.junit.Test;
 import org.tikv.common.KVMockServer.State;
+import org.tikv.kvproto.Metapb;
+import org.tikv.kvproto.Metapb.StoreState;
+import org.tikv.kvproto.Pdpb;
 import org.tikv.raw.RawKVClient;
 
 public class SeekLeaderStoreTest extends MockThreeStoresTest {
@@ -34,12 +38,41 @@ public class SeekLeaderStoreTest extends MockThreeStoresTest {
     RawKVClient client = createClient();
     ByteString key = ByteString.copyFromUtf8("key");
     ByteString value = ByteString.copyFromUtf8("value");
+
     put(key, value);
 
-    client.put(key, value);
     Assert.assertEquals(value, client.get(key).get());
     servers.get(0).setState(State.Fail);
     servers.get(1).setRegion(region.switchPeer(stores.get(1).getId()));
     Assert.assertEquals(value, client.get(key).get());
+
+    remove(key, value);
+  }
+
+  @Test
+  public void testSeekLeaderMeetInvalidStore() {
+    RawKVClient client = createClient();
+    ByteString key = ByteString.copyFromUtf8("key");
+    ByteString value = ByteString.copyFromUtf8("value");
+
+    put(key, value);
+
+    servers.get(0).setState(State.Fail);
+    servers.get(2).setRegion(region.switchPeer(stores.get(2).getId()));
+
+    AtomicInteger i = new AtomicInteger(0);
+    leader.addGetStoreListener(
+        request -> {
+          Metapb.Store.Builder storeBuilder =
+              Metapb.Store.newBuilder().mergeFrom(stores.get((int) request.getStoreId() - 1));
+          if (request.getStoreId() == 0x2 && i.incrementAndGet() > 0) {
+            storeBuilder.setState(StoreState.Tombstone);
+          }
+          return Pdpb.GetStoreResponse.newBuilder().setStore(storeBuilder.build()).build();
+        });
+
+    Assert.assertEquals(value, client.get(key).get());
+
+    remove(key, value);
   }
 }