fix(store): throw on compaction-busy snapshot save; validate data/ on load (#3162)

hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/snapshot/SnapshotHandler.java

@@ -96,7 +96,9 @@ public void onSnapshotSave(final SnapshotWriter writer) throws HgStoreException
             Integer groupId = partitionEngine.getGroupId();
             AtomicInteger state = businessHandler.getState(groupId);
             if (state != null && state.get() == BusinessHandler.doing) {
-                return;
+                throw new HgStoreException(
+                        String.format("Partition %d is busy (compaction in progress), " +
+                                      "snapshot save skipped", groupId));
             }
             // rocks db snapshot
             final String graphSnapshotDir = snapshotDir + File.separator + SNAPSHOT_DATA_PATH;
@@ -172,8 +174,14 @@ public void onSnapshotLoad(final SnapshotReader reader, long committedIndex) thr
 
         // No need to load locally saved snapshots
         if (shouldNotLoad(reader)) {
-            log.info("skip to load snapshot because of should_not_load flag");
-            return;
+            final File dataDir = new File(snapshotDir + File.separator + SNAPSHOT_DATA_PATH);
+            if (dataDir.isDirectory()) {
+                log.info("skip to load snapshot because of should_not_load flag");
+                return;
+            }
+            log.warn("Raft {} should_not_load flag present but data dir {} is missing — " +
+                     "snapshot is corrupt, proceeding to load path",
+                     partitionEngine.getGroupId(), dataDir);
         }
 
         // Use snapshot directly

hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/core/StoreEngineTestBase.java

@@ -28,6 +28,7 @@
 import org.apache.hugegraph.store.meta.Partition;
 import org.apache.hugegraph.store.meta.ShardGroup;
 import org.apache.hugegraph.store.options.HgStoreEngineOptions;
+import org.apache.hugegraph.store.options.JobOptions;
 import org.apache.hugegraph.store.options.RaftRocksdbOptions;
 import org.apache.hugegraph.store.pd.FakePdServiceProvider;
 import org.junit.AfterClass;
@@ -61,6 +62,11 @@ public static void initEngine() {
         options.setGrpcAddress("127.0.0.1:6511");
         options.setRaftAddress("127.0.0.1:6510");
         options.setDataTransfer(new DataManagerImpl());
+        JobOptions jobOptions = new JobOptions();
+        jobOptions.setUninterruptibleCore(2);
+        jobOptions.setUninterruptibleMax(8);
+        jobOptions.setUninterruptibleQueueSize(1024);
+        options.setJobConfig(jobOptions);
 
         options.setFakePdOptions(new HgStoreEngineOptions.FakePdOptions() {{
             setStoreList("127.0.0.1");

hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/core/snapshot/HgSnapshotHandlerTest.java

@@ -18,9 +18,12 @@
 package org.apache.hugegraph.store.core.snapshot;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertThrows;
 
 import java.io.File;
 import java.io.IOException;
+import java.nio.charset.StandardCharsets;
+import java.nio.file.Files;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
@@ -29,8 +32,17 @@
 import org.apache.hugegraph.store.core.StoreEngineTestBase;
 import org.apache.hugegraph.store.meta.Partition;
 import org.apache.hugegraph.store.snapshot.HgSnapshotHandler;
+import org.apache.hugegraph.store.snapshot.SnapshotHandler;
+import org.apache.hugegraph.store.util.HgStoreException;
+
+import com.alipay.sofa.jraft.entity.RaftOutter;
+import com.alipay.sofa.jraft.storage.snapshot.SnapshotReader;
+import com.alipay.sofa.jraft.storage.snapshot.SnapshotWriter;
+import com.google.protobuf.Message;
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
+import org.junit.rules.TemporaryFolder;
 
 import com.alipay.sofa.jraft.entity.RaftOutter;
 import com.alipay.sofa.jraft.storage.snapshot.SnapshotReader;
@@ -42,6 +54,9 @@ public class HgSnapshotHandlerTest extends StoreEngineTestBase {
 
     private static HgSnapshotHandler hgSnapshotHandlerUnderTest;
 
+    @Rule
+    public TemporaryFolder tmpDir = new TemporaryFolder();
+
     @Before
     public void setUp() throws IOException {
         hgSnapshotHandlerUnderTest = new HgSnapshotHandler(createPartitionEngine(0));
@@ -184,4 +199,58 @@ public void testFindFileList() {
 
         // Verify the results
     }
+
+    /**
+     * Test that onSnapshotLoad validates corruption when should_not_load is present but data/ missing.
+     */
+    @Test
+    public void testOnSnapshotLoadFallsThroughWhenShouldNotLoadPresentButDataMissing()
+            throws Exception {
+        // Arrange: snapshot dir has should_not_load but NO data/ subdirectory.
+        File snapDir = tmpDir.newFolder("snapshot-corrupt");
+        File shouldNotLoad = new File(snapDir, "should_not_load");
+        Files.write(shouldNotLoad.toPath(), "saved snapshot".getBytes(StandardCharsets.UTF_8));
+        // data/ deliberately not created
+
+        SnapshotHandler handler = new SnapshotHandler(createPartitionEngine(1));
+        SnapshotReader stubReader = stubReader(snapDir.getAbsolutePath());
+
+        // The fix causes execution to fall through shouldNotLoad() and call
+        // businessHandler.loadSnapshot(missingDataDir) which throws HgStoreException.
+        assertThrows(
+                "onSnapshotLoad must throw when should_not_load present but data/ missing",
+                HgStoreException.class,
+                () -> handler.onSnapshotLoad(stubReader, 0L));
+    }
+
+    /**
+     * Test that onSnapshotLoad skips loading when snapshot is locally saved (both flags present).
+     */
+    @Test
+    public void testOnSnapshotLoadSkipsWhenShouldNotLoadPresentAndDataExists() throws Exception {
+        // Arrange: a healthy local snapshot — both should_not_load and data/ present.
+        File snapDir = tmpDir.newFolder("snapshot-healthy");
+        File shouldNotLoad = new File(snapDir, "should_not_load");
+        Files.write(shouldNotLoad.toPath(), "saved snapshot".getBytes(StandardCharsets.UTF_8));
+        FileUtils.forceMkdir(new File(snapDir, "data"));
+
+        SnapshotHandler handler = new SnapshotHandler(createPartitionEngine(2));
+        SnapshotReader stubReader = stubReader(snapDir.getAbsolutePath());
+
+        // Must not throw — should return early at the should_not_load + data-exists check.
+        handler.onSnapshotLoad(stubReader, 0L);
+    }
+
+    private static SnapshotReader stubReader(String path) {
+        return new SnapshotReader() {
+            @Override public RaftOutter.SnapshotMeta load() { return null; }
+            @Override public String generateURIForCopy() { return null; }
+            @Override public boolean init(Void opts) { return false; }
+            @Override public void shutdown() {}
+            @Override public String getPath() { return path; }
+            @Override public Set<String> listFiles() { return null; }
+            @Override public Message getFileMeta(String fileName) { return null; }
+            @Override public void close() {}
+        };
+    }
 }

hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/core/snapshot/SnapshotHandlerTest.java

@@ -0,0 +1,130 @@
+/*
+ * 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.hugegraph.store.core.snapshot;
+
+import static org.junit.Assert.assertThrows;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.ArgumentMatchers.contains;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+import java.io.File;
+import java.util.Set;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import org.apache.hugegraph.store.HgStoreEngine;
+import org.apache.hugegraph.store.PartitionEngine;
+import org.apache.hugegraph.store.business.BusinessHandler;
+import org.apache.hugegraph.store.snapshot.SnapshotHandler;
+import org.apache.hugegraph.store.util.HgStoreException;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.TemporaryFolder;
+
+import com.alipay.sofa.jraft.entity.RaftOutter;
+import com.alipay.sofa.jraft.storage.snapshot.SnapshotWriter;
+import com.google.protobuf.Message;
+
+public class SnapshotHandlerTest {
+
+    @Rule
+    public TemporaryFolder tmpDir = new TemporaryFolder();
+
+    /**
+     * When state is doing (compaction in progress), onSnapshotSave must throw
+     * immediately. The exception signals jRaft, which will retry the snapshot later.
+     * jRaft's snapshot scheduler runs independently and frequently (default 300s, user config 1800s),
+     * so the next snapshot attempt will succeed after compaction completes.
+     */
+    @Test
+    public void testOnSnapshotSaveThrowsWhenCompactionInProgress() {
+        PartitionEngine mockEngine = mock(PartitionEngine.class);
+        HgStoreEngine mockStoreEngine = mock(HgStoreEngine.class);
+        BusinessHandler mockBusinessHandler = mock(BusinessHandler.class);
+
+        AtomicInteger doingState = new AtomicInteger(BusinessHandler.doing);
+
+        when(mockEngine.getGroupId()).thenReturn(0);
+        when(mockEngine.getStoreEngine()).thenReturn(mockStoreEngine);
+        when(mockStoreEngine.getBusinessHandler()).thenReturn(mockBusinessHandler);
+        when(mockBusinessHandler.getState(0)).thenReturn(doingState);
+
+        SnapshotHandler handler = new SnapshotHandler(mockEngine);
+        SnapshotWriter stubWriter = stubWriter("/tmp/snapshot");
+
+        HgStoreException ex = assertThrows(
+                "onSnapshotSave must throw when state == doing",
+                HgStoreException.class,
+                () -> handler.onSnapshotSave(stubWriter));
+
+        assertTrue("Exception message must mention the partition",
+                   ex.getMessage().contains("0"));
+        assertTrue("Exception message must mention compaction is in progress",
+                   ex.getMessage().contains("compaction in progress"));
+    }
+
+    /**
+     * When state is NOT doing (e.g. compactionDone or null), onSnapshotSave must not throw.
+     * It should proceed and call saveSnapshot with concrete path verification.
+     */
+    @Test
+    public void testOnSnapshotSaveCallsSaveSnapshotWhenNotBusy() throws Exception {
+        PartitionEngine mockEngine = mock(PartitionEngine.class);
+        HgStoreEngine mockStoreEngine = mock(HgStoreEngine.class);
+        BusinessHandler mockBusinessHandler = mock(BusinessHandler.class);
+
+        final String snapshotPath = tmpDir.newFolder("snap-not-busy").getAbsolutePath();
+
+        // state == compactionDone (not doing) — save should proceed immediately
+        AtomicInteger doneState = new AtomicInteger(BusinessHandler.compactionDone);
+
+        when(mockEngine.getGroupId()).thenReturn(0);
+        when(mockEngine.getStoreEngine()).thenReturn(mockStoreEngine);
+        when(mockStoreEngine.getBusinessHandler()).thenReturn(mockBusinessHandler);
+        when(mockBusinessHandler.getState(0)).thenReturn(doneState);
+
+        SnapshotHandler handler = new SnapshotHandler(mockEngine);
+        SnapshotWriter stubWriter = stubWriter(snapshotPath);
+
+        handler.onSnapshotSave(stubWriter);
+
+        // Verify: saveSnapshot was called with concrete path containing expected data dir
+        String expectedDataDir = snapshotPath + File.separator + "data";
+        verify(mockBusinessHandler).saveSnapshot(
+                contains(expectedDataDir),  // Must contain the snapshot path + /data
+                eq(""),                     // graphName (empty string)
+                eq(0));                     // groupId (partition 0)
+    }
+
+    private static SnapshotWriter stubWriter(String path) {
+        return new SnapshotWriter() {
+            @Override public boolean saveMeta(RaftOutter.SnapshotMeta meta) { return false; }
+            @Override public boolean addFile(String fileName, Message fileMeta) { return false; }
+            @Override public boolean removeFile(String fileName) { return false; }
+            @Override public void close(boolean keepDataOnError) {}
+            @Override public boolean init(Void opts) { return false; }
+            @Override public void shutdown() {}
+            @Override public String getPath() { return path; }
+            @Override public Set<String> listFiles() { return null; }
+            @Override public Message getFileMeta(String fileName) { return null; }
+            @Override public void close() {}
+        };
+    }
+}