Examples with Response org.apache.drill.exec.rpc.Response used on opensource projects

Example 11 with Response

use of org.apache.drill.exec.rpc.Response in project drill by axbaretto.

the class ServerAuthenticationHandler method handleSuccess.

private static <S extends ServerConnection<S>, T extends EnumLite> void handleSuccess(final SaslResponseContext<S, T> context, final SaslMessage.Builder challenge, final SaslServer saslServer) throws IOException {
    final S connection = context.connection;
    connection.changeHandlerTo(context.requestHandler);
    connection.finalizeSaslSession();
    // Check the negotiated property before sending the response back to client
    try {
        final String negotiatedQOP = saslServer.getNegotiatedProperty(Sasl.QOP).toString();
        final String expectedQOP = (connection.isEncryptionEnabled()) ? SaslProperties.QualityOfProtection.PRIVACY.getSaslQop() : SaslProperties.QualityOfProtection.AUTHENTICATION.getSaslQop();
        if (!(negotiatedQOP.equals(expectedQOP))) {
            throw new SaslException(String.format("Mismatch in negotiated QOP value: %s and Expected QOP value: %s", negotiatedQOP, expectedQOP));
        }
        // negotiated size of buffer
        if (connection.isEncryptionEnabled()) {
            final int negotiatedRawSendSize = Integer.parseInt(saslServer.getNegotiatedProperty(Sasl.RAW_SEND_SIZE).toString());
            if (negotiatedRawSendSize <= 0) {
                throw new SaslException(String.format("Negotiated rawSendSize: %d is invalid. Please check the configured " + "value of encryption.sasl.max_wrapped_size. It might be configured to a very small value.", negotiatedRawSendSize));
            }
            connection.setWrapSizeLimit(negotiatedRawSendSize);
        }
    } catch (IllegalStateException | NumberFormatException e) {
        throw new SaslException(String.format("Unexpected failure while retrieving negotiated property values (%s)", e.getMessage()), e);
    }
    if (logger.isTraceEnabled()) {
        logger.trace("Authenticated {} successfully using {} from {} with encryption context {}", saslServer.getAuthorizationID(), saslServer.getMechanismName(), connection.getRemoteAddress().toString(), connection.getEncryptionCtxtString());
    }
    // All checks have passed let's send the response back to client before adding handlers.
    context.sender.send(new Response(context.saslResponseType, challenge.build()));
    if (connection.isEncryptionEnabled()) {
        connection.addSecurityHandlers();
    } else {
        // Encryption is not required hence we don't need to hold on to saslServer object.
        connection.disposeSaslServer();
    }
}

Example 12 with Response

use of org.apache.drill.exec.rpc.Response in project drill by axbaretto.

the class CustomHandlerRegistry method handle.

public Response handle(CustomMessage message, DrillBuf dBody) throws RpcException {
    final ParsingHandler<?, ?> handler;
    try (@SuppressWarnings("unused") Closeable 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);
}

Example 13 with Response

use of org.apache.drill.exec.rpc.Response in project drill by axbaretto.

the class DataServerRequestHandler method handle.

@Override
public void handle(DataServerConnection connection, int rpcType, ByteBuf pBody, ByteBuf dBody, ResponseSender sender) throws RpcException {
    assert rpcType == BitData.RpcType.REQ_RECORD_BATCH_VALUE;
    final FragmentRecordBatch fragmentBatch = RpcBus.get(pBody, FragmentRecordBatch.PARSER);
    final AckSender ack = new AckSender(sender);
    // increment so we don't get false returns.
    ack.increment();
    try {
        final IncomingDataBatch batch = new IncomingDataBatch(fragmentBatch, (DrillBuf) dBody, ack);
        final int targetCount = fragmentBatch.getReceivingMinorFragmentIdCount();
        // randomize who gets first transfer (and thus ownership) so memory usage is balanced when we're sharing amongst
        // multiple fragments.
        final int firstOwner = ThreadLocalRandom.current().nextInt(targetCount);
        submit(batch, firstOwner, targetCount);
        submit(batch, 0, firstOwner);
    } catch (IOException | FragmentSetupException e) {
        logger.error("Failure while getting fragment manager. {}", QueryIdHelper.getQueryIdentifiers(fragmentBatch.getQueryId(), fragmentBatch.getReceivingMajorFragmentId(), fragmentBatch.getReceivingMinorFragmentIdList()), e);
        ack.clear();
        sender.send(new Response(BitData.RpcType.ACK, Acks.FAIL));
    } finally {
        // decrement the extra reference we grabbed at the top.
        ack.sendOk();
    }
}

Example 14 with Response

use of org.apache.drill.exec.rpc.Response 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:
            final FragmentStatus status = get(pBody, FragmentStatus.PARSER);
            requestFragmentStatus(status);
            // 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 Ack cancelStatus = requestQueryCancel(queryId);
                if (cancelStatus.getOk()) {
                    sender.send(ControlRpcConfig.OK);
                } else {
                    sender.send(ControlRpcConfig.FAIL);
                }
                break;
            }
        case RpcType.REQ_INITIALIZE_FRAGMENTS_VALUE:
            {
                final InitializeFragments fragments = get(pBody, InitializeFragments.PARSER);
                initializeFragment(fragments);
                sender.send(ControlRpcConfig.OK);
                break;
            }
        case RpcType.REQ_QUERY_STATUS_VALUE:
            {
                final QueryId queryId = get(pBody, QueryId.PARSER);
                final QueryProfile profile = requestQueryStatus(queryId);
                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 15 with Response

use of org.apache.drill.exec.rpc.Response in project drill by apache.

the class ServerAuthenticationHandler method handleSuccess.

private static <S extends ServerConnection<S>, T extends EnumLite> void handleSuccess(final SaslResponseContext<S, T> context, final SaslMessage.Builder challenge, final SaslServer saslServer) throws IOException {
    final S connection = context.connection;
    connection.changeHandlerTo(context.requestHandler);
    connection.finalizeSaslSession();
    // Check the negotiated property before sending the response back to client
    try {
        final String negotiatedQOP = saslServer.getNegotiatedProperty(Sasl.QOP).toString();
        final String expectedQOP = (connection.isEncryptionEnabled()) ? SaslProperties.QualityOfProtection.PRIVACY.getSaslQop() : SaslProperties.QualityOfProtection.AUTHENTICATION.getSaslQop();
        if (!(negotiatedQOP.equals(expectedQOP))) {
            throw new SaslException(String.format("Mismatch in negotiated QOP value: %s and Expected QOP value: %s", negotiatedQOP, expectedQOP));
        }
        // negotiated size of buffer
        if (connection.isEncryptionEnabled()) {
            final int negotiatedRawSendSize = Integer.parseInt(saslServer.getNegotiatedProperty(Sasl.RAW_SEND_SIZE).toString());
            if (negotiatedRawSendSize <= 0) {
                throw new SaslException(String.format("Negotiated rawSendSize: %d is invalid. Please check the configured " + "value of encryption.sasl.max_wrapped_size. It might be configured to a very small value.", negotiatedRawSendSize));
            }
            connection.setWrapSizeLimit(negotiatedRawSendSize);
        }
    } catch (IllegalStateException | NumberFormatException e) {
        throw new SaslException(String.format("Unexpected failure while retrieving negotiated property values (%s)", e.getMessage()), e);
    }
    if (logger.isTraceEnabled()) {
        logger.trace("Authenticated {} successfully using {} from {} with encryption context {}", saslServer.getAuthorizationID(), saslServer.getMechanismName(), connection.getRemoteAddress().toString(), connection.getEncryptionCtxtString());
    }
    // All checks have passed let's send the response back to client before adding handlers.
    context.sender.send(new Response(context.saslResponseType, challenge.build()));
    if (connection.isEncryptionEnabled()) {
        connection.addSecurityHandlers();
    } else {
        // Encryption is not required hence we don't need to hold on to saslServer object.
        connection.disposeSaslServer();
    }
}