fix(store): allocate graph id before batch writes (#3153)
HStore uses a 2-byte GraphId as the graph-isolation boundary. Before this fix, a new graph whose first write used batch PUT or MERGE could encode data with the reserved missing GraphId 0xFFFE. Multiple graphs could then share the same physical RocksDB key range.
---------
Co-authored-by: imbajin <jin@apache.org>
diff --git a/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/business/BusinessHandler.java b/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/business/BusinessHandler.java
index 8133654..58e0056 100644
--- a/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/business/BusinessHandler.java
+++ b/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/business/BusinessHandler.java
@@ -152,6 +152,7 @@
default void doBatch(String graph, int partId, List<BatchEntry> entryList) {
BusinessHandler.TxBuilder builder = txBuilder(graph, partId);
+ BusinessHandler.Tx transaction = builder.build();
try {
for (BatchEntry b : entryList) {
Key start = b.getStartKey();
@@ -185,12 +186,16 @@
}
}
}
- builder.build().commit();
+ transaction.commit();
} catch (Throwable e) {
String msg =
String.format("graph data %s-%s do batch insert with error:", graph, partId);
log.error(msg, e);
- builder.build().rollback();
+ try {
+ transaction.rollback();
+ } catch (Throwable rollbackError) {
+ e.addSuppressed(rollbackError);
+ }
throw e;
}
}
diff --git a/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/business/BusinessHandlerImpl.java b/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/business/BusinessHandlerImpl.java
index 9287bfe..d66093e 100644
--- a/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/business/BusinessHandlerImpl.java
+++ b/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/business/BusinessHandlerImpl.java
@@ -38,6 +38,9 @@
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeoutException;
import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.Lock;
+import java.util.concurrent.locks.ReadWriteLock;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
import java.util.function.BiFunction;
import java.util.function.Consumer;
import java.util.function.Function;
@@ -121,6 +124,8 @@
private static final Map<String, HugeGraphSupplier> GRAPH_SUPPLIER_CACHE =
new ConcurrentHashMap<>();
+ private static final int GRAPH_LOCK_STRIPES = 1024;
+ private static final ReadWriteLock[] GRAPH_LOCKS = createGraphLocks();
private static final int batchSize = 10000;
private static Long indexDataSize = 50 * 1024L;
private static final RocksDBFactory factory = RocksDBFactory.getInstance();
@@ -172,6 +177,20 @@
});
}
+ private static ReadWriteLock[] createGraphLocks() {
+ ReadWriteLock[] locks = new ReadWriteLock[GRAPH_LOCK_STRIPES];
+ for (int i = 0; i < locks.length; i++) {
+ locks[i] = new ReentrantReadWriteLock(true);
+ }
+ return locks;
+ }
+
+ private static ReadWriteLock graphLock(String graph, int partId) {
+ int hash = 31 * partId + graph.hashCode();
+ hash ^= hash >>> 16;
+ return GRAPH_LOCKS[hash & (GRAPH_LOCK_STRIPES - 1)];
+ }
+
public static HugeConfig initRocksdb(Map<String, Object> rocksdbConfig,
RocksdbChangedListener listener) {
// Register rocksdb configuration
@@ -1014,11 +1033,17 @@
public void truncate(String graphName, int partId) throws HgStoreException {
// Each partition corresponds to a rocksdb instance, so the rocksdb instance name is
// rocksdb + partId
- try (RocksDBSession dbSession = getSession(graphName, partId)) {
- dbSession.sessionOp().deleteRange(keyCreator.getStartKey(partId, graphName),
- keyCreator.getEndKey(partId, graphName));
- // Release map ID
- keyCreator.delGraphId(partId, graphName);
+ Lock lifecycleLock = graphLock(graphName, partId).writeLock();
+ lifecycleLock.lock();
+ try {
+ try (RocksDBSession dbSession = getSession(graphName, partId)) {
+ dbSession.sessionOp().deleteRange(keyCreator.getStartKey(partId, graphName),
+ keyCreator.getEndKey(partId, graphName));
+ // Release map ID
+ keyCreator.delGraphId(partId, graphName);
+ }
+ } finally {
+ lifecycleLock.unlock();
}
}
@@ -1287,7 +1312,7 @@
@Override
public TxBuilder txBuilder(String graph, int partId) throws HgStoreException {
- return new TxBuilderImpl(graph, partId, getSession(graph, partId));
+ return new TxBuilderImpl(graph, partId);
}
@Override
@@ -1564,20 +1589,42 @@
private final int partId;
private final RocksDBSession dbSession;
private final SessionOperator op;
+ private final Lock lifecycleLock;
+ private boolean completed;
- private TxBuilderImpl(String graph, int partId, RocksDBSession dbSession) {
+ private TxBuilderImpl(String graph, int partId) {
this.graph = graph;
this.partId = partId;
- this.dbSession = dbSession;
- this.op = this.dbSession.sessionOp();
- this.op.prepare();
+ this.lifecycleLock = graphLock(graph, partId).readLock();
+ this.lifecycleLock.lock();
+
+ RocksDBSession session = null;
+ SessionOperator operator = null;
+ try {
+ session = getSession(graph, partId);
+ operator = session.sessionOp();
+ operator.prepare();
+ } catch (RuntimeException | Error e) {
+ try {
+ if (session != null) {
+ session.close();
+ }
+ } catch (Throwable closeError) {
+ e.addSuppressed(closeError);
+ } finally {
+ this.lifecycleLock.unlock();
+ }
+ throw e;
+ }
+ this.dbSession = session;
+ this.op = operator;
}
@Override
public TxBuilder put(int code, String table, byte[] key, byte[] value) throws
HgStoreException {
try {
- byte[] targetKey = keyCreator.getKey(this.partId, graph, code, key);
+ byte[] targetKey = keyCreator.getKeyOrCreate(this.partId, graph, code, key);
this.op.put(table, targetKey, value);
} catch (DBStoreException e) {
throw new HgStoreException(HgStoreException.EC_RKDB_DOPUT_FAIL, e.toString());
@@ -1642,7 +1689,7 @@
HgStoreException {
try {
- byte[] targetKey = keyCreator.getKey(this.partId, graph, code, key);
+ byte[] targetKey = keyCreator.getKeyOrCreate(this.partId, graph, code, key);
op.merge(table, targetKey, value);
} catch (DBStoreException e) {
throw new HgStoreException(HgStoreException.EC_RKDB_DOMERGE_FAIL, e.toString());
@@ -1655,21 +1702,37 @@
return new Tx() {
@Override
public void commit() throws HgStoreException {
+ if (completed) {
+ return;
+ }
op.commit(); // After an exception occurs in commit, rollback must be
// called, otherwise it will cause the lock not to be released.
- dbSession.close();
+ completed = true;
+ release();
}
@Override
public void rollback() throws HgStoreException {
+ if (completed) {
+ return;
+ }
try {
op.rollback();
} finally {
- dbSession.close();
+ completed = true;
+ release();
}
}
};
}
+
+ private void release() {
+ try {
+ this.dbSession.close();
+ } finally {
+ this.lifecycleLock.unlock();
+ }
+ }
}
public static void clearCache() {
diff --git a/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/business/DataManagerImpl.java b/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/business/DataManagerImpl.java
index 31f111d..3a4b5ac 100644
--- a/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/business/DataManagerImpl.java
+++ b/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/business/DataManagerImpl.java
@@ -222,12 +222,22 @@
@Override
public void write(BatchPutRequest request) {
- BusinessHandler.TxBuilder tx =
+ BusinessHandler.TxBuilder builder =
businessHandler.txBuilder(request.getGraphName(), request.getPartitionId());
- for (BatchPutRequest.KV kv : request.getEntries()) {
- tx.put(kv.getCode(), kv.getTable(), kv.getKey(), kv.getValue());
+ BusinessHandler.Tx transaction = builder.build();
+ try {
+ for (BatchPutRequest.KV kv : request.getEntries()) {
+ builder.put(kv.getCode(), kv.getTable(), kv.getKey(), kv.getValue());
+ }
+ transaction.commit();
+ } catch (Throwable e) {
+ try {
+ transaction.rollback();
+ } catch (Throwable rollbackError) {
+ e.addSuppressed(rollbackError);
+ }
+ throw e;
}
- tx.build().commit();
}
@Override
diff --git a/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/business/DefaultDataMover.java b/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/business/DefaultDataMover.java
index 11f0669..8ba7487 100644
--- a/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/business/DefaultDataMover.java
+++ b/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/business/DefaultDataMover.java
@@ -221,12 +221,22 @@
@Override
public void doWriteData(BatchPutRequest request) {
- BusinessHandler.TxBuilder tx =
+ BusinessHandler.TxBuilder builder =
businessHandler.txBuilder(request.getGraphName(), request.getPartitionId());
- for (BatchPutRequest.KV kv : request.getEntries()) {
- tx.put(kv.getCode(), kv.getTable(), kv.getKey(), kv.getValue());
+ BusinessHandler.Tx transaction = builder.build();
+ try {
+ for (BatchPutRequest.KV kv : request.getEntries()) {
+ builder.put(kv.getCode(), kv.getTable(), kv.getKey(), kv.getValue());
+ }
+ transaction.commit();
+ } catch (Throwable e) {
+ try {
+ transaction.rollback();
+ } catch (Throwable rollbackError) {
+ e.addSuppressed(rollbackError);
+ }
+ throw e;
}
- tx.build().commit();
}
@Override
diff --git a/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/meta/GraphIdManager.java b/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/meta/GraphIdManager.java
index 3a9e2c1..3c95919 100644
--- a/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/meta/GraphIdManager.java
+++ b/hugegraph-store/hg-store-core/src/main/java/org/apache/hugegraph/store/meta/GraphIdManager.java
@@ -17,6 +17,8 @@
package org.apache.hugegraph.store.meta;
+import static org.apache.hugegraph.store.constant.HugeServerTables.VERTEX_TABLE;
+
import java.nio.ByteBuffer;
import java.util.Arrays;
import java.util.List;
@@ -128,8 +130,13 @@
private boolean checkCount(long l) {
var start = new byte[2];
Bits.putShort(start, 0, (short) l);
- try (var itr = sessionBuilder.getSession(partitionId).sessionOp().scan("g+v", start)) {
- return itr == null || !itr.hasNext();
+ try (var session = sessionBuilder.getSession(partitionId)) {
+ if (!session.tableIsExist(VERTEX_TABLE)) {
+ return true;
+ }
+ try (var itr = session.sessionOp().scan(VERTEX_TABLE, start)) {
+ return itr == null || !itr.hasNext();
+ }
}
}
diff --git a/hugegraph-store/hg-store-test/pom.xml b/hugegraph-store/hg-store-test/pom.xml
index 8b7f10d..ed91e01 100644
--- a/hugegraph-store/hg-store-test/pom.xml
+++ b/hugegraph-store/hg-store-test/pom.xml
@@ -244,6 +244,7 @@
</testClassesDirectory>
<includes>
<include>**/CoreSuiteTest.java</include>
+ <include>**/BatchGraphIsolationTest.java</include>
</includes>
</configuration>
</execution>
diff --git a/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/core/BatchGraphIsolationTest.java b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/core/BatchGraphIsolationTest.java
new file mode 100644
index 0000000..c222557
--- /dev/null
+++ b/hugegraph-store/hg-store-test/src/main/java/org/apache/hugegraph/store/core/BatchGraphIsolationTest.java
@@ -0,0 +1,319 @@
+/*
+ * 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;
+
+import static org.apache.hugegraph.store.constant.HugeServerTables.TABLES_MAP;
+import static org.apache.hugegraph.store.constant.HugeServerTables.VERTEX_TABLE;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.nio.ByteOrder;
+import java.nio.charset.StandardCharsets;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
+
+import org.apache.hugegraph.store.UnitTestBase;
+import org.apache.hugegraph.store.business.BusinessHandler;
+import org.apache.hugegraph.store.business.BusinessHandlerImpl;
+import org.apache.hugegraph.store.business.DataManagerImpl;
+import org.apache.hugegraph.store.business.DefaultDataMover;
+import org.apache.hugegraph.store.cmd.request.BatchPutRequest;
+import org.apache.hugegraph.store.grpc.common.Key;
+import org.apache.hugegraph.store.grpc.common.OpType;
+import org.apache.hugegraph.store.grpc.session.BatchEntry;
+import org.apache.hugegraph.store.meta.PartitionManager;
+import org.apache.hugegraph.store.options.HgStoreEngineOptions;
+import org.apache.hugegraph.store.options.RaftRocksdbOptions;
+import org.apache.hugegraph.store.pd.FakePdServiceProvider;
+import org.apache.hugegraph.store.pd.PdProvider;
+import org.junit.AfterClass;
+import org.junit.Assert;
+import org.junit.BeforeClass;
+import org.junit.Test;
+import org.mockito.Mockito;
+
+import com.alipay.sofa.jraft.util.StorageOptionsFactory;
+import com.google.protobuf.ByteString;
+
+public class BatchGraphIsolationTest {
+
+ private static final int PARTITION_ID = 0;
+ private static final int EMPTY_PARTITION_ID = 1;
+ private static final int KEY_CODE = 0;
+ private static final int EMPTY_PARTITION_KEY_CODE = 32768;
+ private static final byte[] SHARED_KEY =
+ "shared-key".getBytes(StandardCharsets.UTF_8);
+
+ private static Path databasePath;
+ private static BusinessHandler handler;
+
+ @BeforeClass
+ public static void setup() throws IOException {
+ databasePath = Files.createTempDirectory("hugegraph-batch-graph-isolation-");
+
+ Map<String, Object> rocksdbConfig = new HashMap<>();
+ rocksdbConfig.put("rocksdb.write_buffer_size", "1048576");
+ StorageOptionsFactory.releaseAllOptions();
+ RaftRocksdbOptions.initRocksdbGlobalConfig(rocksdbConfig);
+ BusinessHandlerImpl.initRocksdb(rocksdbConfig, null);
+
+ HgStoreEngineOptions options = new HgStoreEngineOptions();
+ options.setDataPath(databasePath.toString());
+ options.setRaftPath(databasePath.toString());
+
+ HgStoreEngineOptions.FakePdOptions fakePdOptions =
+ new HgStoreEngineOptions.FakePdOptions();
+ fakePdOptions.setPartitionCount(2);
+ fakePdOptions.setPeersList("127.0.0.1");
+ fakePdOptions.setStoreList("127.0.0.1");
+ options.setFakePdOptions(fakePdOptions);
+
+ PdProvider pdProvider = new FakePdServiceProvider(fakePdOptions);
+ PartitionManager partitionManager = new PartitionManager(pdProvider, options) {
+
+ @Override
+ public String getDbDataPath(int partitionId, String dbName) {
+ return databasePath.resolve("data").toString();
+ }
+
+ @Override
+ public boolean hasPartition(String graphName, int partitionId) {
+ return partitionId == PARTITION_ID || partitionId == EMPTY_PARTITION_ID;
+ }
+
+ @Override
+ public List<Integer> getLeaderPartitionIds(String graph) {
+ return Collections.singletonList(PARTITION_ID);
+ }
+ };
+ handler = new BusinessHandlerImpl(partitionManager);
+ handler.createTable("setup", PARTITION_ID, VERTEX_TABLE);
+ }
+
+ @AfterClass
+ public static void teardown() {
+ if (handler != null) {
+ handler.closeAll();
+ }
+ if (databasePath != null) {
+ UnitTestBase.deleteDir(databasePath.toFile());
+ }
+ }
+
+ @Test(timeout = 5000L)
+ public void testFirstBatchOnEmptyPartitionCompletes() {
+ String graph = "first-batch-empty-partition";
+ byte[] value = "first-value".getBytes(StandardCharsets.UTF_8);
+
+ Assert.assertFalse(handler.existsTable(graph, EMPTY_PARTITION_ID, VERTEX_TABLE));
+ writeBatch(graph, EMPTY_PARTITION_ID, EMPTY_PARTITION_KEY_CODE,
+ OpType.OP_TYPE_PUT, value);
+ Assert.assertArrayEquals(value, readByCode(graph, EMPTY_PARTITION_KEY_CODE));
+ }
+
+ @Test(timeout = 10000L)
+ public void testTruncateWaitsForInFlightBatch() throws Exception {
+ String graph = "in-flight-batch-graph";
+ String nextGraph = "graph-after-truncate";
+ byte[] pendingValue = "pending-value".getBytes(StandardCharsets.UTF_8);
+ byte[] nextValue = "next-value".getBytes(StandardCharsets.UTF_8);
+ BusinessHandler.TxBuilder builder = handler.txBuilder(graph, PARTITION_ID);
+ BusinessHandler.Tx transaction = builder.build();
+ ExecutorService executor = Executors.newSingleThreadExecutor();
+ CountDownLatch truncateStarted = new CountDownLatch(1);
+ boolean committed = false;
+
+ try {
+ put(builder, VERTEX_TABLE, pendingValue);
+ Future<?> truncate = executor.submit(() -> {
+ truncateStarted.countDown();
+ handler.truncate(graph, PARTITION_ID);
+ });
+
+ Assert.assertTrue(truncateStarted.await(1L, TimeUnit.SECONDS));
+ try {
+ truncate.get(500L, TimeUnit.MILLISECONDS);
+ Assert.fail("truncate must wait for the in-flight batch");
+ } catch (TimeoutException expected) {
+ // Expected until the transaction releases its graph ID reservation.
+ }
+
+ transaction.commit();
+ committed = true;
+ truncate.get(5L, TimeUnit.SECONDS);
+
+ writeBatch(nextGraph, OpType.OP_TYPE_PUT, nextValue);
+ Assert.assertArrayEquals(nextValue, read(nextGraph));
+ } finally {
+ if (!committed) {
+ transaction.rollback();
+ }
+ executor.shutdownNow();
+ }
+ }
+
+ @Test
+ public void testDataManagerRollsBackFailedBatch() {
+ BusinessHandler mockHandler = Mockito.mock(BusinessHandler.class);
+ BusinessHandler.TxBuilder mockBuilder = Mockito.mock(BusinessHandler.TxBuilder.class);
+ BusinessHandler.Tx mockTransaction = Mockito.mock(BusinessHandler.Tx.class);
+ BatchPutRequest request = batchPutRequest();
+ BatchPutRequest.KV entry = request.getEntries().get(0);
+ RuntimeException failure = new RuntimeException("injected put failure");
+ Mockito.when(mockHandler.txBuilder(request.getGraphName(), request.getPartitionId()))
+ .thenReturn(mockBuilder);
+ Mockito.when(mockBuilder.build()).thenReturn(mockTransaction);
+ Mockito.doThrow(failure).when(mockBuilder)
+ .put(entry.getCode(), entry.getTable(), entry.getKey(), entry.getValue());
+ DataManagerImpl dataManager = new DataManagerImpl();
+ dataManager.setBusinessHandler(mockHandler);
+
+ assertWriteFails(failure, () -> dataManager.write(request));
+
+ Mockito.verify(mockTransaction).rollback();
+ }
+
+ @SuppressWarnings("deprecation")
+ @Test
+ public void testDefaultDataMoverRollsBackFailedBatch() {
+ BusinessHandler mockHandler = Mockito.mock(BusinessHandler.class);
+ BusinessHandler.TxBuilder mockBuilder = Mockito.mock(BusinessHandler.TxBuilder.class);
+ BusinessHandler.Tx mockTransaction = Mockito.mock(BusinessHandler.Tx.class);
+ BatchPutRequest request = batchPutRequest();
+ BatchPutRequest.KV entry = request.getEntries().get(0);
+ RuntimeException failure = new RuntimeException("injected put failure");
+ Mockito.when(mockHandler.txBuilder(request.getGraphName(), request.getPartitionId()))
+ .thenReturn(mockBuilder);
+ Mockito.when(mockBuilder.build()).thenReturn(mockTransaction);
+ Mockito.doThrow(failure).when(mockBuilder)
+ .put(entry.getCode(), entry.getTable(), entry.getKey(), entry.getValue());
+ DefaultDataMover dataMover = new DefaultDataMover();
+ dataMover.setBusinessHandler(mockHandler);
+
+ assertWriteFails(failure, () -> dataMover.doWriteData(request));
+
+ Mockito.verify(mockTransaction).rollback();
+ }
+
+ @Test
+ public void testBatchPutKeepsGraphsIsolated() {
+ String graph1 = "batch-put-graph-1";
+ String graph2 = "batch-put-graph-2";
+ byte[] value1 = "graph-1-value".getBytes(StandardCharsets.UTF_8);
+ byte[] value2 = "graph-2-value".getBytes(StandardCharsets.UTF_8);
+
+ writeBatch(graph1, OpType.OP_TYPE_PUT, value1);
+ writeBatch(graph2, OpType.OP_TYPE_PUT, value2);
+
+ Assert.assertArrayEquals(value1, read(graph1));
+ Assert.assertArrayEquals(value2, read(graph2));
+
+ handler.truncate(graph2, PARTITION_ID);
+ Assert.assertNull(read(graph2));
+ Assert.assertArrayEquals(value1, read(graph1));
+ }
+
+ @Test
+ public void testBatchMergeKeepsGraphsIsolated() {
+ String graph1 = "batch-merge-graph-1";
+ String graph2 = "batch-merge-graph-2";
+
+ writeBatch(graph1, OpType.OP_TYPE_MERGE, longToBytes(10L));
+ writeBatch(graph2, OpType.OP_TYPE_MERGE, longToBytes(20L));
+
+ Assert.assertEquals(10L, bytesToLong(read(graph1)));
+ Assert.assertEquals(20L, bytesToLong(read(graph2)));
+ }
+
+ private static void writeBatch(String graph, OpType type, byte[] value) {
+ writeBatch(graph, PARTITION_ID, type, value);
+ }
+
+ private static void writeBatch(String graph, int partitionId, OpType type, byte[] value) {
+ writeBatch(graph, partitionId, KEY_CODE, type, value);
+ }
+
+ private static void writeBatch(String graph, int partitionId, int code, OpType type,
+ byte[] value) {
+ Key key = Key.newBuilder()
+ .setCode(code)
+ .setKey(ByteString.copyFrom(SHARED_KEY))
+ .build();
+ BatchEntry entry = BatchEntry.newBuilder()
+ .setOpType(type)
+ .setTable(TABLES_MAP.get(VERTEX_TABLE))
+ .setStartKey(key)
+ .setValue(ByteString.copyFrom(value))
+ .build();
+ handler.doBatch(graph, partitionId, Collections.singletonList(entry));
+ }
+
+ private static void put(BusinessHandler.TxBuilder builder, String table, byte[] value) {
+ builder.put(KEY_CODE, table, SHARED_KEY, value);
+ }
+
+ private static byte[] read(String graph) {
+ return readByCode(graph, KEY_CODE);
+ }
+
+ private static byte[] readByCode(String graph, int code) {
+ return handler.doGet(graph, code, VERTEX_TABLE, SHARED_KEY);
+ }
+
+ private static BatchPutRequest batchPutRequest() {
+ BatchPutRequest request = new BatchPutRequest();
+ request.setGraphName("failed-transfer-graph");
+ request.setPartitionId(PARTITION_ID);
+ request.getEntries()
+ .add(BatchPutRequest.KV.of(VERTEX_TABLE, KEY_CODE, SHARED_KEY,
+ "value".getBytes(StandardCharsets.UTF_8)));
+ return request;
+ }
+
+ private static void assertWriteFails(RuntimeException expected, Runnable writer) {
+ try {
+ writer.run();
+ Assert.fail("batch write must propagate its failure");
+ } catch (RuntimeException actual) {
+ Assert.assertSame(expected, actual);
+ }
+ }
+
+ private static byte[] longToBytes(long value) {
+ return ByteBuffer.allocate(Long.BYTES)
+ .order(ByteOrder.LITTLE_ENDIAN)
+ .putLong(value)
+ .array();
+ }
+
+ private static long bytesToLong(byte[] value) {
+ return ByteBuffer.wrap(value)
+ .order(ByteOrder.LITTLE_ENDIAN)
+ .getLong();
+ }
+}
