fix can't retry InstallSnapshot problem( #605 ) (#606)

* (fix) Fixed the problem that all InstallSnapshot sessions will be blocked when a new InstallSnapshot request arrives, see #605
sofastack · Jun 16, 2021 · a1fa8c3 · a1fa8c3
1 parent 3817c8d
commit a1fa8c3
Show file tree

Hide file tree

Showing 2 changed files with 109 additions and 2 deletions.
diff --git a/jraft-core/src/main/java/com/alipay/sofa/jraft/storage/snapshot/SnapshotExecutorImpl.java b/jraft-core/src/main/java/com/alipay/sofa/jraft/storage/snapshot/SnapshotExecutorImpl.java
@@ -611,7 +611,7 @@ boolean registerDownloadingSnapshot(final DownloadingSnapshot ds) {
                 // this RPC.
                 saved = m;
                 this.downloadingSnapshot.set(ds);
-                result = false;
+                result = true;
             } else if (m.request.getMeta().getLastIncludedIndex() > ds.request.getMeta().getLastIncludedIndex()) {
                 // |is| is older
                 LOG.warn("Register DownloadingSnapshot failed: is installing a newer one, lastIncludeIndex={}.",
@@ -648,7 +648,7 @@ boolean registerDownloadingSnapshot(final DownloadingSnapshot ds) {
         }
         if (saved != null) {
             // Respond replaced session
-            LOG.warn("Register DownloadingSnapshot failed: interrupted by retry installling request.");
+            LOG.warn("Register DownloadingSnapshot failed: interrupted by retry installing request.");
             saved.done.sendResponse(RpcFactoryHelper //
                 .responseFactory() //
                 .newResponse(InstallSnapshotResponse.getDefaultInstance(), RaftError.EINTR,

diff --git a/jraft-core/src/test/java/com/alipay/sofa/jraft/storage/SnapshotExecutorTest.java b/jraft-core/src/test/java/com/alipay/sofa/jraft/storage/SnapshotExecutorTest.java
@@ -17,14 +17,20 @@
 package com.alipay.sofa.jraft.storage;
 
 import java.nio.ByteBuffer;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.Future;
+import java.util.concurrent.atomic.AtomicInteger;
 
+import com.alipay.sofa.jraft.error.RaftError;
+import com.alipay.sofa.jraft.test.MockAsyncContext;
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.ArgumentCaptor;
 import org.mockito.Mock;
 import org.mockito.Mockito;
+import org.mockito.invocation.InvocationOnMock;
 import org.mockito.runners.MockitoJUnitRunner;
 
 import com.alipay.sofa.jraft.FSMCaller;
@@ -58,6 +64,7 @@
 import com.alipay.sofa.jraft.util.Utils;
 import com.google.protobuf.ByteString;
 import com.google.protobuf.Message;
+import org.mockito.stubbing.Answer;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotNull;
@@ -135,6 +142,106 @@ public void teardown() throws Exception {
         this.timerManager.shutdown();
     }
 
+    @Test
+    public void testRetryInstallSnapshot() throws Exception {
+        final RpcRequests.InstallSnapshotRequest.Builder irb = RpcRequests.InstallSnapshotRequest.newBuilder();
+        irb.setGroupId("test");
+        irb.setPeerId(this.addr.toString());
+        irb.setServerId("localhost:8080");
+        irb.setUri("remote://localhost:8080/99");
+        irb.setTerm(0);
+        irb.setMeta(RaftOutter.SnapshotMeta.newBuilder().setLastIncludedIndex(1).setLastIncludedTerm(2));
+
+        Mockito.when(this.raftClientService.connect(new Endpoint("localhost", 8080))).thenReturn(true);
+
+        final FutureImpl<Message> future = new FutureImpl<>();
+        final RpcRequests.GetFileRequest.Builder rb = RpcRequests.GetFileRequest.newBuilder().setReaderId(99)
+            .setFilename(Snapshot.JRAFT_SNAPSHOT_META_FILE).setCount(Integer.MAX_VALUE).setOffset(0)
+            .setReadPartly(true);
+
+        //mock get metadata
+        ArgumentCaptor<RpcResponseClosure> argument = ArgumentCaptor.forClass(RpcResponseClosure.class);
+
+        final CountDownLatch retryLatch = new CountDownLatch(1);
+        final CountDownLatch answerLatch = new CountDownLatch(1);
+        Mockito.when(
+            this.raftClientService.getFile(eq(new Endpoint("localhost", 8080)), eq(rb.build()),
+                eq(this.copyOpts.getTimeoutMs()), argument.capture())).thenAnswer(new Answer<Future<Message>>() {
+            AtomicInteger count = new AtomicInteger(0);
+
+            @Override
+            public Future<Message> answer(InvocationOnMock invocation) throws Throwable {
+                if (count.incrementAndGet() == 1) {
+                    retryLatch.countDown();
+                    answerLatch.await();
+                    Thread.sleep(1000);
+                    return future;
+                } else {
+                    throw new IllegalStateException("shouldn't be called more than once");
+                }
+            }
+        });
+
+        final MockAsyncContext installContext = new MockAsyncContext();
+        final MockAsyncContext retryInstallContext = new MockAsyncContext();
+        Utils.runInThread(new Runnable() {
+            @Override
+            public void run() {
+                SnapshotExecutorTest.this.executor.installSnapshot(irb.build(),
+                    RpcRequests.InstallSnapshotResponse.newBuilder(), new RpcRequestClosure(installContext));
+            }
+        });
+
+        Thread.sleep(500);
+        retryLatch.await();
+        Utils.runInThread(new Runnable() {
+            @Override
+            public void run() {
+                answerLatch.countDown();
+                SnapshotExecutorTest.this.executor.installSnapshot(irb.build(),
+                    RpcRequests.InstallSnapshotResponse.newBuilder(), new RpcRequestClosure(retryInstallContext));
+            }
+        });
+
+        RpcResponseClosure<RpcRequests.GetFileResponse> closure = argument.getValue();
+        final ByteBuffer metaBuf = this.table.saveToByteBufferAsRemote();
+        closure.setResponse(RpcRequests.GetFileResponse.newBuilder().setReadSize(metaBuf.remaining()).setEof(true)
+            .setData(ByteString.copyFrom(metaBuf)).build());
+
+        //mock get file
+        argument = ArgumentCaptor.forClass(RpcResponseClosure.class);
+        rb.setFilename("testFile");
+        rb.setCount(this.raftOptions.getMaxByteCountPerRpc());
+        Mockito.when(
+            this.raftClientService.getFile(eq(new Endpoint("localhost", 8080)), eq(rb.build()),
+                eq(this.copyOpts.getTimeoutMs()), argument.capture())).thenReturn(future);
+
+        closure.run(Status.OK());
+        Thread.sleep(500);
+        closure = argument.getValue();
+        closure.setResponse(RpcRequests.GetFileResponse.newBuilder().setReadSize(100).setEof(true)
+            .setData(ByteString.copyFrom(new byte[100])).build());
+
+        final ArgumentCaptor<LoadSnapshotClosure> loadSnapshotArg = ArgumentCaptor.forClass(LoadSnapshotClosure.class);
+        Mockito.when(this.fSMCaller.onSnapshotLoad(loadSnapshotArg.capture())).thenReturn(true);
+        closure.run(Status.OK());
+        Thread.sleep(2000);
+        final LoadSnapshotClosure done = loadSnapshotArg.getValue();
+        final SnapshotReader reader = done.start();
+        assertNotNull(reader);
+        assertEquals(1, reader.listFiles().size());
+        assertTrue(reader.listFiles().contains("testFile"));
+        done.run(Status.OK());
+        this.executor.join();
+        assertEquals(2, this.executor.getLastSnapshotTerm());
+        assertEquals(1, this.executor.getLastSnapshotIndex());
+        assertNotNull(installContext.getResponseObject());
+        assertNotNull(retryInstallContext.getResponseObject());
+        assertEquals(installContext.as(RpcRequests.ErrorResponse.class).getErrorCode(), RaftError.EINTR.getNumber());
+        assertTrue(retryInstallContext.as(RpcRequests.InstallSnapshotResponse.class).hasSuccess());
+
+    }
+
     @Test
     public void testInstallSnapshot() throws Exception {
         final RpcRequests.InstallSnapshotRequest.Builder irb = RpcRequests.InstallSnapshotRequest.newBuilder();