@Test public void testRemoteAccess() throws TableNotExistException, InterruptedException, NetworkException { final TableComponents<String, Integer, Integer> tableComponents = tables.getTableComponents(TABLE_ID); final KVUSerializer<String, Integer, Integer> kvuSerializer = tableComponents.getSerializer(); final StreamingCodec<String> keyCodec = kvuSerializer.getKeyCodec(); final StreamingCodec<Integer> valueCodec = kvuSerializer.getValueCodec(); final ResettingCountDownLatch replyLatch = new ResettingCountDownLatch(1); final AtomicReference<Integer> reply = new AtomicReference<>(); doAnswer(invocation -> { final Object[] arguments = invocation.getArguments(); final String destId = (String) arguments[0]; final long opId = (long) arguments[1]; final DataValue dataValue = (DataValue) arguments[3]; final Integer value = dataValue == null ? null : valueCodec.decode(dataValue.getValue().array()); final boolean isSuccess = (boolean) arguments[4]; if (!isSuccess || !destId.equals(ORIG_EXECUTOR_ID)) { fail(String.format("Operation with id %d is failed", opId)); } reply.set(value); replyLatch.countDown(); LOG.log(Level.FINE, "Access result. opId: {0}, value: {1}, destId: {2}", new Object[]{opId, value, destId}); return null; }).when(mockMsgSender).sendTableAccessResMsg(anyString(), anyLong(), anyString(), anyObject(), anyBoolean()); final String key = "key"; final Integer value = 1; final Pair<DataKey, DataValue> dataPair = getDataPair(key, value, keyCodec, valueCodec); final AtomicInteger opIdCounter = new AtomicInteger(0); final TableAccessReqMsg putMsg = TableAccessReqMsg.newBuilder() .setOrigId(ORIG_EXECUTOR_ID) .setOpType(OpType.PUT) .setTableId(TABLE_ID) .setDataKey(dataPair.getKey()) .setDataValue(dataPair.getValue()) .build(); remoteAccessOpHandler.onTableAccessReqMsg(opIdCounter.getAndIncrement(), putMsg); replyLatch.awaitAndReset(1); assertEquals(null, reply.get()); final TableAccessReqMsg getMsg = TableAccessReqMsg.newBuilder() .setOrigId(ORIG_EXECUTOR_ID) .setOpType(OpType.GET) .setTableId(TABLE_ID) .setDataKey(dataPair.getKey()) .setDataValue(null) .build(); remoteAccessOpHandler.onTableAccessReqMsg(opIdCounter.getAndIncrement(), getMsg); replyLatch.awaitAndReset(1); assertEquals(value, reply.get()); final TableAccessReqMsg updateMsg = TableAccessReqMsg.newBuilder() .setOrigId(ORIG_EXECUTOR_ID) .setOpType(OpType.UPDATE) .setTableId(TABLE_ID) .setDataKey(dataPair.getKey()) .setDataValue(dataPair.getValue()) .build(); remoteAccessOpHandler.onTableAccessReqMsg(opIdCounter.getAndIncrement(), updateMsg); replyLatch.awaitAndReset(1); assertEquals(Integer.valueOf(value + value), reply.get()); final TableAccessReqMsg putIfAbsentMsg = TableAccessReqMsg.newBuilder() .setOrigId(ORIG_EXECUTOR_ID) .setOpType(OpType.PUT_IF_ABSENT) .setTableId(TABLE_ID) .setDataKey(dataPair.getKey()) .setDataValue(dataPair.getValue()) .build(); remoteAccessOpHandler.onTableAccessReqMsg(opIdCounter.getAndIncrement(), putIfAbsentMsg); replyLatch.awaitAndReset(1); assertEquals(Integer.valueOf(value + value), reply.get()); remoteAccessOpHandler.onTableAccessReqMsg(opIdCounter.getAndIncrement(), getMsg); replyLatch.awaitAndReset(1); assertEquals(Integer.valueOf(value + value), reply.get()); final TableAccessReqMsg removeMsg = TableAccessReqMsg.newBuilder() .setOrigId(ORIG_EXECUTOR_ID) .setOpType(OpType.REMOVE) .setTableId(TABLE_ID) .setDataKey(dataPair.getKey()) .setDataValue(null) .build(); remoteAccessOpHandler.onTableAccessReqMsg(opIdCounter.getAndIncrement(), removeMsg); replyLatch.awaitAndReset(1); assertEquals(Integer.valueOf(value + value), reply.get()); remoteAccessOpHandler.onTableAccessReqMsg(opIdCounter.getAndIncrement(), getMsg); replyLatch.awaitAndReset(1); assertEquals(null, reply.get()); verify(mockMsgSender, times(7)) .sendTableAccessResMsg(anyString(), anyLong(), anyString(), anyObject(), anyBoolean()); }
<K, V, U> void onTableAccessReqMsg(final long opId, final TableAccessReqMsg msg) { final String origEvalId = msg.getOrigId(); final OpType opType = msg.getOpType(); final boolean replyRequired = msg.getReplyRequired(); final String tableId = msg.getTableId(); final TableComponents<K, V, U> tableComponents; try { tableComponents = tablesFuture.get().getTableComponents(tableId); } catch (TableNotExistException e) { redirectToMaster(opId, msg, origEvalId); return; } if (!registerOp(tableId, tableComponents)) { LOG.log(Level.INFO, "Operations for table {0} are being flushed. No more ops will be accepted.", tableId); redirectToMaster(opId, msg, origEvalId); return; } final KVUSerializer<K, V, U> kvuSerializer = tableComponents.getSerializer(); final StreamingCodec<K> keyCodec = kvuSerializer.getKeyCodec(); final StreamingCodec<V> valueCodec = kvuSerializer.getValueCodec(); final Codec<U> updateValueCodec = kvuSerializer.getUpdateValueCodec(); final BlockPartitioner<K> blockPartitioner = tableComponents.getBlockPartitioner(); if (msg.getIsSingleKey()) { final DataKey dataKey = msg.getDataKey(); final DataValue dataValue = msg.getDataValue(); final K decodedKey = keyCodec.decode(dataKey.getKey().array()); final V decodedValue = opType.equals(OpType.PUT) || opType.equals(OpType.PUT_IF_ABSENT) ? valueCodec.decode(dataValue.getValue().array()) : null; final U decodedUpdateValue = opType.equals(OpType.UPDATE) ? updateValueCodec.decode(dataValue.getValue().array()) : null; final int blockId = blockPartitioner.getBlockId(decodedKey); final SingleKeyDataOpMetadata<K, V, U> operation = new SingleKeyDataOpMetadata<>(origEvalId, opId, opType, replyRequired, tableId, blockId, decodedKey, decodedValue, decodedUpdateValue); final int threadIdx = getThreadIdx(blockId); handlerThreads.get(threadIdx).enqueue(operation); } else { final DataKeys dataKeys = msg.getDataKeys(); final DataValues dataValues = msg.getDataValues(); final List<K> keyList = new ArrayList<>(); dataKeys.getKeys().forEach(key -> keyList.add(keyCodec.decode(key.array()))); final List<V> valueList; final List<U> updateValueList; switch (opType) { case PUT: valueList = new ArrayList<>(); updateValueList = Collections.emptyList(); dataValues.getValues().forEach(value -> valueList.add(valueCodec.decode(value.array()))); break; case GET: case GET_OR_INIT: valueList = Collections.emptyList(); updateValueList = Collections.emptyList(); break; case UPDATE: valueList = Collections.emptyList(); updateValueList = new ArrayList<>(); dataValues.getValues().forEach(value -> updateValueList.add(updateValueCodec.decode(value.array()))); break; default: throw new RuntimeException("Undefined type of OpMetadata"); } final int blockId = blockPartitioner.getBlockId(keyList.get(0)); final MultiKeyDataOpMetadata<K, V, ?> operation = new MultiKeyDataOpMetadata<>(origEvalId, opId, opType, replyRequired, tableId, blockId, keyList, valueList, updateValueList); final int threadIdx = getThreadIdx(blockId); handlerThreads.get(threadIdx).enqueue(operation); } }
RemoteAccessOpHandler { <K, V, U> void onTableAccessReqMsg(final long opId, final TableAccessReqMsg msg) { final String origEvalId = msg.getOrigId(); final OpType opType = msg.getOpType(); final boolean replyRequired = msg.getReplyRequired(); final String tableId = msg.getTableId(); final TableComponents<K, V, U> tableComponents; try { tableComponents = tablesFuture.get().getTableComponents(tableId); } catch (TableNotExistException e) { redirectToMaster(opId, msg, origEvalId); return; } if (!registerOp(tableId, tableComponents)) { LOG.log(Level.INFO, "Operations for table {0} are being flushed. No more ops will be accepted.", tableId); redirectToMaster(opId, msg, origEvalId); return; } final KVUSerializer<K, V, U> kvuSerializer = tableComponents.getSerializer(); final StreamingCodec<K> keyCodec = kvuSerializer.getKeyCodec(); final StreamingCodec<V> valueCodec = kvuSerializer.getValueCodec(); final Codec<U> updateValueCodec = kvuSerializer.getUpdateValueCodec(); final BlockPartitioner<K> blockPartitioner = tableComponents.getBlockPartitioner(); if (msg.getIsSingleKey()) { final DataKey dataKey = msg.getDataKey(); final DataValue dataValue = msg.getDataValue(); final K decodedKey = keyCodec.decode(dataKey.getKey().array()); final V decodedValue = opType.equals(OpType.PUT) || opType.equals(OpType.PUT_IF_ABSENT) ? valueCodec.decode(dataValue.getValue().array()) : null; final U decodedUpdateValue = opType.equals(OpType.UPDATE) ? updateValueCodec.decode(dataValue.getValue().array()) : null; final int blockId = blockPartitioner.getBlockId(decodedKey); final SingleKeyDataOpMetadata<K, V, U> operation = new SingleKeyDataOpMetadata<>(origEvalId, opId, opType, replyRequired, tableId, blockId, decodedKey, decodedValue, decodedUpdateValue); final int threadIdx = getThreadIdx(blockId); handlerThreads.get(threadIdx).enqueue(operation); } else { final DataKeys dataKeys = msg.getDataKeys(); final DataValues dataValues = msg.getDataValues(); final List<K> keyList = new ArrayList<>(); dataKeys.getKeys().forEach(key -> keyList.add(keyCodec.decode(key.array()))); final List<V> valueList; final List<U> updateValueList; switch (opType) { case PUT: valueList = new ArrayList<>(); updateValueList = Collections.emptyList(); dataValues.getValues().forEach(value -> valueList.add(valueCodec.decode(value.array()))); break; case GET: case GET_OR_INIT: valueList = Collections.emptyList(); updateValueList = Collections.emptyList(); break; case UPDATE: valueList = Collections.emptyList(); updateValueList = new ArrayList<>(); dataValues.getValues().forEach(value -> updateValueList.add(updateValueCodec.decode(value.array()))); break; default: throw new RuntimeException("Undefined type of OpMetadata"); } final int blockId = blockPartitioner.getBlockId(keyList.get(0)); final MultiKeyDataOpMetadata<K, V, ?> operation = new MultiKeyDataOpMetadata<>(origEvalId, opId, opType, replyRequired, tableId, blockId, keyList, valueList, updateValueList); final int threadIdx = getThreadIdx(blockId); handlerThreads.get(threadIdx).enqueue(operation); } } }
RemoteAccessOpHandler { <K, V, U> void onTableAccessReqMsg(final long opId, final TableAccessReqMsg msg) { final String origEvalId = msg.getOrigId(); final OpType opType = msg.getOpType(); final boolean replyRequired = msg.getReplyRequired(); final String tableId = msg.getTableId(); final TableComponents<K, V, U> tableComponents; try { tableComponents = tablesFuture.get().getTableComponents(tableId); } catch (TableNotExistException e) { redirectToMaster(opId, msg, origEvalId); return; } if (!registerOp(tableId, tableComponents)) { LOG.log(Level.INFO, "Operations for table {0} are being flushed. No more ops will be accepted.", tableId); redirectToMaster(opId, msg, origEvalId); return; } final KVUSerializer<K, V, U> kvuSerializer = tableComponents.getSerializer(); final StreamingCodec<K> keyCodec = kvuSerializer.getKeyCodec(); final StreamingCodec<V> valueCodec = kvuSerializer.getValueCodec(); final Codec<U> updateValueCodec = kvuSerializer.getUpdateValueCodec(); final BlockPartitioner<K> blockPartitioner = tableComponents.getBlockPartitioner(); if (msg.getIsSingleKey()) { final DataKey dataKey = msg.getDataKey(); final DataValue dataValue = msg.getDataValue(); final K decodedKey = keyCodec.decode(dataKey.getKey().array()); final V decodedValue = opType.equals(OpType.PUT) || opType.equals(OpType.PUT_IF_ABSENT) ? valueCodec.decode(dataValue.getValue().array()) : null; final U decodedUpdateValue = opType.equals(OpType.UPDATE) ? updateValueCodec.decode(dataValue.getValue().array()) : null; final int blockId = blockPartitioner.getBlockId(decodedKey); final SingleKeyDataOpMetadata<K, V, U> operation = new SingleKeyDataOpMetadata<>(origEvalId, opId, opType, replyRequired, tableId, blockId, decodedKey, decodedValue, decodedUpdateValue); final int threadIdx = getThreadIdx(blockId); handlerThreads.get(threadIdx).enqueue(operation); } else { final DataKeys dataKeys = msg.getDataKeys(); final DataValues dataValues = msg.getDataValues(); final List<K> keyList = new ArrayList<>(); dataKeys.getKeys().forEach(key -> keyList.add(keyCodec.decode(key.array()))); final List<V> valueList; final List<U> updateValueList; switch (opType) { case PUT: valueList = new ArrayList<>(); updateValueList = Collections.emptyList(); dataValues.getValues().forEach(value -> valueList.add(valueCodec.decode(value.array()))); break; case GET: case GET_OR_INIT: valueList = Collections.emptyList(); updateValueList = Collections.emptyList(); break; case UPDATE: valueList = Collections.emptyList(); updateValueList = new ArrayList<>(); dataValues.getValues().forEach(value -> updateValueList.add(updateValueCodec.decode(value.array()))); break; default: throw new RuntimeException("Undefined type of OpMetadata"); } final int blockId = blockPartitioner.getBlockId(keyList.get(0)); final MultiKeyDataOpMetadata<K, V, ?> operation = new MultiKeyDataOpMetadata<>(origEvalId, opId, opType, replyRequired, tableId, blockId, keyList, valueList, updateValueList); final int threadIdx = getThreadIdx(blockId); handlerThreads.get(threadIdx).enqueue(operation); } } @Inject private  RemoteAccessOpHandler(final InjectionFuture<Tables> tablesFuture,
                                @Parameter(DriverIdentifier.class) final String driverId,
                                @Parameter(ExecutorIdentifier.class) final String executorId,
                                @Parameter(HandlerQueueSize.class) final int queueSize,
                                @Parameter(NumRemoteOpsHandlerThreads.class) final int numHandlerThreads,
                                final InjectionFuture<MessageSender> msgSenderFuture); }
RemoteAccessOpHandler { <K, V, U> void onTableAccessReqMsg(final long opId, final TableAccessReqMsg msg) { final String origEvalId = msg.getOrigId(); final OpType opType = msg.getOpType(); final boolean replyRequired = msg.getReplyRequired(); final String tableId = msg.getTableId(); final TableComponents<K, V, U> tableComponents; try { tableComponents = tablesFuture.get().getTableComponents(tableId); } catch (TableNotExistException e) { redirectToMaster(opId, msg, origEvalId); return; } if (!registerOp(tableId, tableComponents)) { LOG.log(Level.INFO, "Operations for table {0} are being flushed. No more ops will be accepted.", tableId); redirectToMaster(opId, msg, origEvalId); return; } final KVUSerializer<K, V, U> kvuSerializer = tableComponents.getSerializer(); final StreamingCodec<K> keyCodec = kvuSerializer.getKeyCodec(); final StreamingCodec<V> valueCodec = kvuSerializer.getValueCodec(); final Codec<U> updateValueCodec = kvuSerializer.getUpdateValueCodec(); final BlockPartitioner<K> blockPartitioner = tableComponents.getBlockPartitioner(); if (msg.getIsSingleKey()) { final DataKey dataKey = msg.getDataKey(); final DataValue dataValue = msg.getDataValue(); final K decodedKey = keyCodec.decode(dataKey.getKey().array()); final V decodedValue = opType.equals(OpType.PUT) || opType.equals(OpType.PUT_IF_ABSENT) ? valueCodec.decode(dataValue.getValue().array()) : null; final U decodedUpdateValue = opType.equals(OpType.UPDATE) ? updateValueCodec.decode(dataValue.getValue().array()) : null; final int blockId = blockPartitioner.getBlockId(decodedKey); final SingleKeyDataOpMetadata<K, V, U> operation = new SingleKeyDataOpMetadata<>(origEvalId, opId, opType, replyRequired, tableId, blockId, decodedKey, decodedValue, decodedUpdateValue); final int threadIdx = getThreadIdx(blockId); handlerThreads.get(threadIdx).enqueue(operation); } else { final DataKeys dataKeys = msg.getDataKeys(); final DataValues dataValues = msg.getDataValues(); final List<K> keyList = new ArrayList<>(); dataKeys.getKeys().forEach(key -> keyList.add(keyCodec.decode(key.array()))); final List<V> valueList; final List<U> updateValueList; switch (opType) { case PUT: valueList = new ArrayList<>(); updateValueList = Collections.emptyList(); dataValues.getValues().forEach(value -> valueList.add(valueCodec.decode(value.array()))); break; case GET: case GET_OR_INIT: valueList = Collections.emptyList(); updateValueList = Collections.emptyList(); break; case UPDATE: valueList = Collections.emptyList(); updateValueList = new ArrayList<>(); dataValues.getValues().forEach(value -> updateValueList.add(updateValueCodec.decode(value.array()))); break; default: throw new RuntimeException("Undefined type of OpMetadata"); } final int blockId = blockPartitioner.getBlockId(keyList.get(0)); final MultiKeyDataOpMetadata<K, V, ?> operation = new MultiKeyDataOpMetadata<>(origEvalId, opId, opType, replyRequired, tableId, blockId, keyList, valueList, updateValueList); final int threadIdx = getThreadIdx(blockId); handlerThreads.get(threadIdx).enqueue(operation); } } @Inject private  RemoteAccessOpHandler(final InjectionFuture<Tables> tablesFuture,
                                @Parameter(DriverIdentifier.class) final String driverId,
                                @Parameter(ExecutorIdentifier.class) final String executorId,
                                @Parameter(HandlerQueueSize.class) final int queueSize,
                                @Parameter(NumRemoteOpsHandlerThreads.class) final int numHandlerThreads,
                                final InjectionFuture<MessageSender> msgSenderFuture);  }
RemoteAccessOpHandler { <K, V, U> void onTableAccessReqMsg(final long opId, final TableAccessReqMsg msg) { final String origEvalId = msg.getOrigId(); final OpType opType = msg.getOpType(); final boolean replyRequired = msg.getReplyRequired(); final String tableId = msg.getTableId(); final TableComponents<K, V, U> tableComponents; try { tableComponents = tablesFuture.get().getTableComponents(tableId); } catch (TableNotExistException e) { redirectToMaster(opId, msg, origEvalId); return; } if (!registerOp(tableId, tableComponents)) { LOG.log(Level.INFO, "Operations for table {0} are being flushed. No more ops will be accepted.", tableId); redirectToMaster(opId, msg, origEvalId); return; } final KVUSerializer<K, V, U> kvuSerializer = tableComponents.getSerializer(); final StreamingCodec<K> keyCodec = kvuSerializer.getKeyCodec(); final StreamingCodec<V> valueCodec = kvuSerializer.getValueCodec(); final Codec<U> updateValueCodec = kvuSerializer.getUpdateValueCodec(); final BlockPartitioner<K> blockPartitioner = tableComponents.getBlockPartitioner(); if (msg.getIsSingleKey()) { final DataKey dataKey = msg.getDataKey(); final DataValue dataValue = msg.getDataValue(); final K decodedKey = keyCodec.decode(dataKey.getKey().array()); final V decodedValue = opType.equals(OpType.PUT) || opType.equals(OpType.PUT_IF_ABSENT) ? valueCodec.decode(dataValue.getValue().array()) : null; final U decodedUpdateValue = opType.equals(OpType.UPDATE) ? updateValueCodec.decode(dataValue.getValue().array()) : null; final int blockId = blockPartitioner.getBlockId(decodedKey); final SingleKeyDataOpMetadata<K, V, U> operation = new SingleKeyDataOpMetadata<>(origEvalId, opId, opType, replyRequired, tableId, blockId, decodedKey, decodedValue, decodedUpdateValue); final int threadIdx = getThreadIdx(blockId); handlerThreads.get(threadIdx).enqueue(operation); } else { final DataKeys dataKeys = msg.getDataKeys(); final DataValues dataValues = msg.getDataValues(); final List<K> keyList = new ArrayList<>(); dataKeys.getKeys().forEach(key -> keyList.add(keyCodec.decode(key.array()))); final List<V> valueList; final List<U> updateValueList; switch (opType) { case PUT: valueList = new ArrayList<>(); updateValueList = Collections.emptyList(); dataValues.getValues().forEach(value -> valueList.add(valueCodec.decode(value.array()))); break; case GET: case GET_OR_INIT: valueList = Collections.emptyList(); updateValueList = Collections.emptyList(); break; case UPDATE: valueList = Collections.emptyList(); updateValueList = new ArrayList<>(); dataValues.getValues().forEach(value -> updateValueList.add(updateValueCodec.decode(value.array()))); break; default: throw new RuntimeException("Undefined type of OpMetadata"); } final int blockId = blockPartitioner.getBlockId(keyList.get(0)); final MultiKeyDataOpMetadata<K, V, ?> operation = new MultiKeyDataOpMetadata<>(origEvalId, opId, opType, replyRequired, tableId, blockId, keyList, valueList, updateValueList); final int threadIdx = getThreadIdx(blockId); handlerThreads.get(threadIdx).enqueue(operation); } } @Inject private  RemoteAccessOpHandler(final InjectionFuture<Tables> tablesFuture,
                                @Parameter(DriverIdentifier.class) final String driverId,
                                @Parameter(ExecutorIdentifier.class) final String executorId,
                                @Parameter(HandlerQueueSize.class) final int queueSize,
                                @Parameter(NumRemoteOpsHandlerThreads.class) final int numHandlerThreads,
                                final InjectionFuture<MessageSender> msgSenderFuture);   }