Examples with CustomMessage org.apache.drill.exec.proto.BitControl.CustomMessage used on opensource projects

Example 1 with CustomMessage

use of org.apache.drill.exec.proto.BitControl.CustomMessage in project drill by apache.

the class ControlMessageHandler method handle.

@Override
public void handle(ControlConnection connection, int rpcType, ByteBuf pBody, ByteBuf dBody, ResponseSender sender) throws RpcException {
    if (RpcConstants.EXTRA_DEBUGGING) {
        logger.debug("Received bit com message of type {}", rpcType);
    }
    switch(rpcType) {
        case RpcType.REQ_CANCEL_FRAGMENT_VALUE:
            {
                final FragmentHandle handle = get(pBody, FragmentHandle.PARSER);
                cancelFragment(handle);
                sender.send(ControlRpcConfig.OK);
                break;
            }
        case RpcType.REQ_CUSTOM_VALUE:
            {
                final CustomMessage customMessage = get(pBody, CustomMessage.PARSER);
                sender.send(handlerRegistry.handle(customMessage, (DrillBuf) dBody));
                break;
            }
        case RpcType.REQ_RECEIVER_FINISHED_VALUE:
            {
                final FinishedReceiver finishedReceiver = get(pBody, FinishedReceiver.PARSER);
                receivingFragmentFinished(finishedReceiver);
                sender.send(ControlRpcConfig.OK);
                break;
            }
        case RpcType.REQ_FRAGMENT_STATUS_VALUE:
            bee.getContext().getWorkBus().statusUpdate(get(pBody, FragmentStatus.PARSER));
            // TODO: Support a type of message that has no response.
            sender.send(ControlRpcConfig.OK);
            break;
        case RpcType.REQ_QUERY_CANCEL_VALUE:
            {
                final QueryId queryId = get(pBody, QueryId.PARSER);
                final Foreman foreman = bee.getForemanForQueryId(queryId);
                if (foreman != null) {
                    foreman.cancel();
                    sender.send(ControlRpcConfig.OK);
                } else {
                    sender.send(ControlRpcConfig.FAIL);
                }
                break;
            }
        case RpcType.REQ_INITIALIZE_FRAGMENTS_VALUE:
            {
                final InitializeFragments fragments = get(pBody, InitializeFragments.PARSER);
                for (int i = 0; i < fragments.getFragmentCount(); i++) {
                    startNewRemoteFragment(fragments.getFragment(i));
                }
                sender.send(ControlRpcConfig.OK);
                break;
            }
        case RpcType.REQ_QUERY_STATUS_VALUE:
            {
                final QueryId queryId = get(pBody, QueryId.PARSER);
                final Foreman foreman = bee.getForemanForQueryId(queryId);
                if (foreman == null) {
                    throw new RpcException("Query not running on node.");
                }
                final QueryProfile profile = foreman.getQueryManager().getQueryProfile();
                sender.send(new Response(RpcType.RESP_QUERY_STATUS, profile));
                break;
            }
        case RpcType.REQ_UNPAUSE_FRAGMENT_VALUE:
            {
                final FragmentHandle handle = get(pBody, FragmentHandle.PARSER);
                resumeFragment(handle);
                sender.send(ControlRpcConfig.OK);
                break;
            }
        default:
            throw new RpcException("Not yet supported.");
    }
}

Example 2 with CustomMessage

use of org.apache.drill.exec.proto.BitControl.CustomMessage in project drill by apache.

the class CustomHandlerRegistry method handle.

public Response handle(CustomMessage message, DrillBuf dBody) throws RpcException {
    final ParsingHandler<?, ?> handler;
    try (AutoCloseableLock lock = read.open()) {
        handler = handlers.get(message.getType());
    }
    if (handler == null) {
        throw new UserRpcException(endpoint, "Unable to handle message.", new IllegalStateException(String.format("Unable to handle message. The message type provided [%d] did not have a registered handler.", message.getType())));
    }
    final CustomResponse<?> customResponse = handler.onMessage(message.getMessage(), dBody);
    @SuppressWarnings("unchecked") final CustomMessage responseMessage = CustomMessage.newBuilder().setMessage(ByteString.copyFrom(((Controller.CustomSerDe<Object>) handler.getResponseSerDe()).serializeToSend(customResponse.getMessage()))).setType(message.getType()).build();
    // make sure we don't pass in a null array.
    final ByteBuf[] dBodies = customResponse.getBodies() == null ? new DrillBuf[0] : customResponse.getBodies();
    return new Response(RpcType.RESP_CUSTOM, responseMessage, dBodies);
}