org.elasticsearch.transport.TransportRequest Java Examples

@Override
public Connection openConnection(DiscoveryNode node, ConnectionProfile profile) {
    return new CloseableConnection() {
        @Override
        public DiscoveryNode getNode() {
            return node;
        }

        @Override
        public void sendRequest(long requestId, String action, TransportRequest request, TransportRequestOptions options)
            throws TransportException {
            requests.put(requestId, Tuple.tuple(node, action));
            onSendRequest(requestId, action, request, node);
        }

        @Override
        public String toString() {
            return "MockTransportConnection{node=" + node + ", requests=" + requests + '}';
        }
    };
}
 

@Before
public void createServices() {
    localNode = newDiscoveryNode("local");
    otherNode1 = newDiscoveryNode("other1");
    otherNode2 = newDiscoveryNode("other2");

    deterministicTaskQueue = new DeterministicTaskQueue(builder().put(NODE_NAME_SETTING.getKey(), "node").build(), random());

    final MockTransport transport = new MockTransport() {
        @Override
        protected void onSendRequest(long requestId, String action, TransportRequest request, DiscoveryNode node) {
            throw new AssertionError("unexpected " + action);
        }
    };

    transportService = transport.createTransportService(
        Settings.EMPTY,
        deterministicTaskQueue.getThreadPool(),
        TransportService.NOOP_TRANSPORT_INTERCEPTOR,
        boundTransportAddress -> localNode,
        null
    );
}
 

@Override
public void logSSLException(TransportRequest request, Throwable t, String action, Task task) {

    if(!checkTransportFilter(Category.SSL_EXCEPTION, action, getUser(), request)) {
        return;
    }

    final TransportAddress remoteAddress = getRemoteAddress();

    final List<AuditMessage> msgs = RequestResolver.resolve(Category.SSL_EXCEPTION, Origin.TRANSPORT, action, null, getUser(), false, null, remoteAddress, request,
            getThreadContextHeaders(), task, resolver, clusterService, settings, logRequestBody, resolveIndices, resolveBulkRequests, opendistrosecurityIndex, excludeSensitiveHeaders, t);

    for(AuditMessage msg: msgs) {
        save(msg);
    }
}
 

@Override
public Connection openConnection(DiscoveryNode node, ConnectionProfile profile) {
    final Optional<DisruptableMockTransport> matchingTransport = getDisruptableMockTransport(node.getAddress());
    if (matchingTransport.isPresent()) {
        return new CloseableConnection() {
            @Override
            public DiscoveryNode getNode() {
                return node;
            }

            @Override
            public void sendRequest(long requestId, String action, TransportRequest request, TransportRequestOptions options)
                throws TransportException {
                onSendRequest(requestId, action, request, matchingTransport.get());
            }

            @Override
            public String toString() {
                return "DisruptableMockTransportConnection{node=" + node + '}';
            }
        };
    } else {
        throw new ConnectTransportException(node, "node " + node + " does not exist");
    }
}
 

public <TRequest extends TransportRequest, TResponse extends TransportResponse> void sendRequest(
    String action,
    String node,
    TRequest request,
    ActionListener<TResponse> listener,
    TransportResponseHandler<TResponse> handler,
    TransportRequestOptions options) {

    DiscoveryNode discoveryNode = clusterService.state().nodes().get(node);
    if (discoveryNode == null) {
        listener.onFailure(new NodeNotConnectedException(null,
            String.format(Locale.ENGLISH, "node \"%s\" not found in cluster state!", node)));
        return;
    }
    transportService.sendRequest(discoveryNode, action, request, options, handler);
}
 

@Test
public void testPutHeadChunkRunnableFileDoesntGrow() throws Exception {
    // this test is rather slow, tune wait time in PutHeadChunkRunnable?
    expectedException.expect(HeadChunkFileTooSmallException.class);

    File file = File.createTempFile("test", "");
    File notExisting = new File("./does/not/exist");
    try (final FileOutputStream outputStream = new FileOutputStream(file)) {
        outputStream.write(new byte[]{0x65});
    }
    UUID transferId = UUID.randomUUID();
    BlobTransferTarget transferTarget = mock(BlobTransferTarget.class);
    TransportService transportService = mock(TransportService.class);
    DiscoveryNode discoveryNode = mock(DiscoveryNode.class);

    DigestBlob digestBlob = mock(DigestBlob.class);
    when(digestBlob.file()).thenReturn(notExisting);
    when(digestBlob.getContainerFile()).thenReturn(file);
    PutHeadChunkRunnable runnable = new PutHeadChunkRunnable(
        digestBlob, 5, transportService, transferTarget, discoveryNode, transferId
    );

    @SuppressWarnings("unchecked")
    TransportFuture<TransportResponse.Empty> result = mock(TransportFuture.class);
    when(transportService.submitRequest(
        eq(discoveryNode),
        eq(BlobHeadRequestHandler.Actions.PUT_BLOB_HEAD_CHUNK),
        any(TransportRequest.class),
        any(TransportRequestOptions.class),
        eq(EmptyTransportResponseHandler.INSTANCE_SAME)
    )).thenReturn(result);

    runnable.run();
    verify(digestBlob).getContainerFile();
}
 

@Override
public List<TransportInterceptor> getTransportInterceptors(NamedWriteableRegistry namedWriteableRegistry) {
    return Arrays.asList(new TransportInterceptor() {
        @Override
        public <T extends TransportRequest> TransportRequestHandler<T> interceptHandler(String action, String executor,
                                                                                        boolean forceExecution,
                                                                                        TransportRequestHandler<T> actualHandler) {
            if (TransportService.isValidActionName(action) == false) {
                throw new IllegalArgumentException("invalid action name [" + action + "] must start with one of: " +
                    TransportService.VALID_ACTION_PREFIXES );
            }
            return actualHandler;
        }
    });
}
 

@Override
public void logSucceededLogin(String effectiveUser, boolean securityadmin, String initiatingUser, TransportRequest request, String action, Task task) {

    if(!checkTransportFilter(Category.AUTHENTICATED, action, effectiveUser, request)) {
        return;
    }

    final TransportAddress remoteAddress = getRemoteAddress();
    final List<AuditMessage> msgs = RequestResolver.resolve(Category.AUTHENTICATED, getOrigin(), action, null, effectiveUser, securityadmin, initiatingUser,remoteAddress, request, getThreadContextHeaders(), task, resolver, clusterService, settings, logRequestBody, resolveIndices, resolveBulkRequests, opendistrosecurityIndex, excludeSensitiveHeaders, null);

    for(AuditMessage msg: msgs) {
        save(msg);
    }
}
 

@Test
public void testOnFailureOnListenerIsCalledIfNodeIsNotInClusterState() throws Exception {
    Transports transports = new Transports(clusterService, mock(TransportService.class));
    ActionListener actionListener = mock(ActionListener.class);
    transports.sendRequest("actionName",
        "invalid", mock(TransportRequest.class), actionListener, mock(TransportResponseHandler.class));

    verify(actionListener, times(1)).onFailure(any(ConnectTransportException.class));
}
 

public <TRequest extends TransportRequest, TResponse extends TransportResponse> void sendRequest(
        String actionName,
        String node,
        TRequest request,
        ActionListener<TResponse> listener,
        TransportResponseHandler<TResponse> transportResponseHandler) {
    DiscoveryNode discoveryNode = clusterService.state().nodes().get(node);
    if (discoveryNode == null) {
        listener.onFailure(new IllegalArgumentException(
            String.format(Locale.ENGLISH, "node \"%s\" not found in cluster state!", node)));
        return;
    }
    transportService.sendRequest(discoveryNode, actionName, request, transportResponseHandler);
}
 

@Override
public <Request extends TransportRequest> void registerRequestHandler(RequestHandlerRegistry<Request> reg) {
    synchronized (requestHandlerMutex) {
        if (requestHandlers.containsKey(reg.getAction())) {
            throw new IllegalArgumentException("transport handlers for action " + reg.getAction() + " is already registered");
        }
        HashMap<String, RequestHandlerRegistry> copy = new HashMap<>(requestHandlers);
        copy.put(reg.getAction(), reg);
        requestHandlers = Map.copyOf(copy);
    }
}
 

@Override
public void logFailedLogin(String effectiveUser, boolean securityadmin, String initiatingUser, TransportRequest request, Task task) {
    final String action = null;

    if(!checkTransportFilter(Category.FAILED_LOGIN, action, effectiveUser, request)) {
        return;
    }

    final TransportAddress remoteAddress = getRemoteAddress();
    final List<AuditMessage> msgs = RequestResolver.resolve(Category.FAILED_LOGIN, getOrigin(), action, null, effectiveUser, securityadmin, initiatingUser, remoteAddress, request, getThreadContextHeaders(), task, resolver, clusterService, settings, logRequestBody, resolveIndices, resolveBulkRequests, opendistrosecurityIndex, excludeSensitiveHeaders, null);

    for(AuditMessage msg: msgs) {
        save(msg);
    }
}
 

@Override
public void logMissingPrivileges(String privilege, TransportRequest request, Task task) {
    final String action = null;

    if(!checkTransportFilter(Category.MISSING_PRIVILEGES, privilege, getUser(), request)) {
        return;
    }

    final TransportAddress remoteAddress = getRemoteAddress();
    final List<AuditMessage> msgs = RequestResolver.resolve(Category.MISSING_PRIVILEGES, getOrigin(), action, privilege, getUser(), null, null, remoteAddress, request, getThreadContextHeaders(), task, resolver, clusterService, settings, logRequestBody, resolveIndices, resolveBulkRequests, opendistrosecurityIndex, excludeSensitiveHeaders, null);

    for(AuditMessage msg: msgs) {
        save(msg);
    }
}
 

@Override
public void logGrantedPrivileges(String privilege, TransportRequest request, Task task) {
    final String action = null;

    if(!checkTransportFilter(Category.GRANTED_PRIVILEGES, privilege, getUser(), request)) {
        return;
    }

    final TransportAddress remoteAddress = getRemoteAddress();
    final List<AuditMessage> msgs = RequestResolver.resolve(Category.GRANTED_PRIVILEGES, getOrigin(), action, privilege, getUser(), null, null, remoteAddress, request, getThreadContextHeaders(), task, resolver, clusterService, settings, logRequestBody, resolveIndices, resolveBulkRequests, opendistrosecurityIndex, excludeSensitiveHeaders, null);

    for(AuditMessage msg: msgs) {
        save(msg);
    }
}
 

@Override
public void logBadHeaders(TransportRequest request, String action, Task task) {

    if(!checkTransportFilter(Category.BAD_HEADERS, action, getUser(), request)) {
        return;
    }

    final TransportAddress remoteAddress = getRemoteAddress();
    final List<AuditMessage> msgs = RequestResolver.resolve(Category.BAD_HEADERS, getOrigin(), action, null, getUser(), null, null, remoteAddress, request, getThreadContextHeaders(), task, resolver, clusterService, settings, logRequestBody, resolveIndices, resolveBulkRequests, opendistrosecurityIndex, excludeSensitiveHeaders, null);

    for(AuditMessage msg: msgs) {
        save(msg);
    }
}
 

@Override
public void sendRequest(long requestId, String action, TransportRequest request, TransportRequestOptions options)
    throws IOException, TransportException {
    TransportAddress address = connection.getNode().getAddress();
    SendRequestBehavior behavior = sendBehaviors.getOrDefault(address, defaultSendRequest);
    if (behavior == null) {
        connection.sendRequest(requestId, action, request, options);
    } else {
        behavior.sendRequest(connection, requestId, action, request, options);
    }
}
 

/**
 * Adds a rule that will cause ignores each send request, simulating an unresponsive node
 * and failing to connect once the rule was added.
 */
public void addUnresponsiveRule(TransportAddress transportAddress) {
    transport().addConnectBehavior(transportAddress, (transport, discoveryNode, profile) -> {
        throw new ConnectTransportException(discoveryNode, "UNRESPONSIVE: simulated");
    });

    transport().addSendBehavior(
        transportAddress,
        new StubbableTransport.SendRequestBehavior() {
            private Set<Transport.Connection> toClose = ConcurrentHashMap.newKeySet();

            @Override
            public void sendRequest(Transport.Connection connection,
                                    long requestId,
                                    String action,
                                    TransportRequest request,
                                    TransportRequestOptions options) {
                // don't send anything, the receiving node is unresponsive
                toClose.add(connection);
            }

            @Override
            public void clearCallback() {
                // close to simulate that tcp-ip eventually times out and closes connection
                // (necessary to ensure transport eventually responds).
                try {
                    IOUtils.close(toClose);
                } catch (IOException e) {
                    throw new RuntimeException(e);
                }
            }
        }
    );
}
 

protected void onSendRequest(long requestId, String action, TransportRequest request,
                             DisruptableMockTransport destinationTransport) {

    assert destinationTransport.getLocalNode().equals(getLocalNode()) == false :
        "non-local message from " + getLocalNode() + " to itself";

    destinationTransport.execute(action, new Runnable() {
        @Override
        public void run() {
            final ConnectionStatus connectionStatus = getConnectionStatus(destinationTransport.getLocalNode());
            switch (connectionStatus) {
                case BLACK_HOLE:
                case BLACK_HOLE_REQUESTS_ONLY:
                    onBlackholedDuringSend(requestId, action, destinationTransport);
                    break;

                case DISCONNECTED:
                    onDisconnectedDuringSend(requestId, action, destinationTransport);
                    break;

                case CONNECTED:
                    onConnectedDuringSend(requestId, action, request, destinationTransport);
                    break;

                default:
                    throw new AssertionError("unexpected status: " + connectionStatus);
            }
        }

        @Override
        public String toString() {
            return getRequestDescription(requestId, action, destinationTransport.getLocalNode());
        }
    });
}
 

/**
 * Registers a task without parent task
 * <p>
 * Returns the task manager tracked task or null if the task doesn't support the task manager
 */
public Task register(String type, String action, TransportRequest request) {
    Task task = request.createTask(taskIdGenerator.incrementAndGet(), type, action);
    if (task != null) {
        if (logger.isTraceEnabled()) {
            logger.trace("register {} [{}] [{}] [{}]", task.getId(), type, action, task.getDescription());
        }

        if (task instanceof CancellableTask) {
            CancellableTask cancellableTask = (CancellableTask) task;
            CancellableTaskHolder holder = new CancellableTaskHolder(cancellableTask);
            CancellableTaskHolder oldHolder = cancellableTasks.put(task.getId(), holder);
            assert oldHolder == null;
            // Check if this task was banned before we start it
            if (task.getParentTaskId().isSet() && banedParents.isEmpty() == false) {
                String reason = banedParents.get(task.getParentTaskId());
                if (reason != null) {
                    try {
                        holder.cancel(reason);
                        throw new IllegalStateException("Task cancelled before it started: " + reason);
                    } finally {
                        // let's clean up the registration
                        unregister(task);
                    }
                }
            }
        } else {
            Task previousTask = tasks.put(task.getId(), task);
            assert previousTask == null;
        }

    }
    return task;
}
 

protected ShardFetchRequest(TransportRequest originalRequest, long id, IntArrayList list, ScoreDoc lastEmittedDoc) {
    super(originalRequest);
    this.id = id;
    this.docIds = list.buffer;
    this.size = list.size();
    this.lastEmittedDoc = lastEmittedDoc;
}
 

public <TRequest extends TransportRequest, TResponse extends TransportResponse> void sendRequest(
    String action,
    String node,
    TRequest request,
    ActionListener<TResponse> listener,
    TransportResponseHandler<TResponse> handler) {
    sendRequest(action, node, request, listener, handler, TransportRequestOptions.EMPTY);
}
 

/**
 * Convenience to turn
 *
 *      `action.execute(request, listener) -> void`
 *
 *  into
 *
 *      `execute(action, request, convertResp) -> resultFuture`
 */
public static <Req extends TransportRequest, Resp extends TransportResponse, Result> CompletableFuture<Result> execute(
    TransportAction<Req, Resp> transportAction,
    Req request,
    Function<? super Resp, Result> convertResponse) {

    FutureActionListener<Resp, Result> listener = new FutureActionListener<>(convertResponse);
    transportAction.execute(request, listener);
    return listener;
}
 

@Override
public List<ActionHandler<? extends TransportRequest, ? extends TransportResponse>> getActions() {
    return Arrays.asList(
        new ActionHandler<>(PutChunkAction.INSTANCE, TransportPutChunkAction.class),
        new ActionHandler<>(StartBlobAction.INSTANCE, TransportStartBlobAction.class),
        new ActionHandler<>(DeleteBlobAction.INSTANCE, TransportDeleteBlobAction.class));
}
 

@Override
public final <Request extends TransportRequest, Response extends TransportResponse, RequestBuilder extends ActionRequestBuilder<Request, Response, RequestBuilder>> ActionFuture<Response> execute(
        Action<Request, Response, RequestBuilder> action, Request request) {
    PlainActionFuture<Response> actionFuture = PlainActionFuture.newFuture();
    execute(action, request, actionFuture);
    return actionFuture;
}
 

private void performAction(final DiscoveryNode node, final String action, final boolean isPrimaryAction,
                           final TransportRequest requestToPerform) {
    transportService.sendRequest(node, action, requestToPerform, transportOptions, new TransportResponseHandler<Response>() {

        @Override
        public Response read(StreamInput in) throws IOException {
            return TransportReplicationAction.this.read(in);
        }

        @Override
        public String executor() {
            return ThreadPool.Names.SAME;
        }

        @Override
        public void handleResponse(Response response) {
            finishOnSuccess(response);
        }

        @Override
        public void handleException(TransportException exp) {
            try {
                // if we got disconnected from the node, or the node / shard is not in the right state (being closed)
                final Throwable cause = exp.unwrapCause();
                if (cause instanceof ConnectTransportException || cause instanceof NodeClosedException ||
                    (isPrimaryAction && retryPrimaryException(cause))) {
                    logger.trace(() -> new ParameterizedMessage(
                            "received an error from node [{}] for request [{}], scheduling a retry",
                            node.getId(), requestToPerform), exp);
                    retry(exp);
                } else {
                    finishAsFailed(exp);
                }
            } catch (Exception e) {
                e.addSuppressed(exp);
                finishWithUnexpectedFailure(e);
            }
        }
    });
}
 

@Override
public <T extends TransportRequest> TransportRequestHandler<T> interceptHandler(String action, String executor,
                                                                                boolean forceExecution,
                                                                                TransportRequestHandler<T> actualHandler) {
    for (TransportInterceptor interceptor : this.transportInterceptors) {
        actualHandler = interceptor.interceptHandler(action, executor, forceExecution, actualHandler);
    }
    return actualHandler;
}
 

protected void waitForNewMasterAndRetry(String actionName,
                                        ClusterStateObserver observer,
                                        TransportRequest request,
                                        ActionListener<Void> listener,
                                        Predicate<ClusterState> changePredicate) {
    observer.waitForNextChange(new ClusterStateObserver.Listener() {
        @Override
        public void onNewClusterState(ClusterState state) {
            if (LOGGER.isTraceEnabled()) {
                LOGGER.trace("new cluster state [{}] after waiting for master election for shard entry [{}]", state, request);
            }
            sendShardAction(actionName, state, request, listener);
        }

        @Override
        public void onClusterServiceClose() {
            LOGGER.warn("node closed while execution action [{}] for shard entry [{}]", actionName, request);
            listener.onFailure(new NodeClosedException(clusterService.localNode()));
        }

        @Override
        public void onTimeout(TimeValue timeout) {
            // we wait indefinitely for a new master
            assert false;
        }
    }, changePredicate);
}
 

@Override
public <Request extends TransportRequest,
        Response extends TransportResponse,
        RequestBuilder extends ActionRequestBuilder<Request, Response, RequestBuilder>> void doExecute(
    Action<Request, Response, RequestBuilder> action,
    Request request,
    ActionListener<Response> listener) {

    // Discard the task because the Client interface doesn't use it.
    executeLocally(action, request, listener);
}
 

/**
 * Execute an {@link Action} locally, returning that {@link Task} used to track it, and linking an {@link ActionListener}. Prefer this
 * method if you don't need access to the task when listening for the response. This is the method used to implement the {@link Client}
 * interface.
 */
public <Request extends TransportRequest,
        Response extends TransportResponse> Task executeLocally(GenericAction<Request, Response> action,
                                                                Request request,
                                                                ActionListener<Response> listener) {
    return transportAction(action).execute(request, listener);
}
 

private void sendShardAction(final String actionName, final ClusterState currentState, final TransportRequest request, final ActionListener<Void> listener) {
    ClusterStateObserver observer = new ClusterStateObserver(currentState, clusterService, null, LOGGER);
    DiscoveryNode masterNode = currentState.nodes().getMasterNode();
    Predicate<ClusterState> changePredicate = MasterNodeChangePredicate.build(currentState);
    if (masterNode == null) {
        LOGGER.warn("no master known for action [{}] for shard entry [{}]", actionName, request);
        waitForNewMasterAndRetry(actionName, observer, request, listener, changePredicate);
    } else {
        LOGGER.debug("sending [{}] to [{}] for shard entry [{}]", actionName, masterNode.getId(), request);
        transportService.sendRequest(masterNode,
            actionName, request, new EmptyTransportResponseHandler(ThreadPool.Names.SAME) {
                @Override
                public void handleResponse(TransportResponse.Empty response) {
                    listener.onResponse(null);
                }

                @Override
                public void handleException(TransportException exp) {
                    if (isMasterChannelException(exp)) {
                        waitForNewMasterAndRetry(actionName, observer, request, listener, changePredicate);
                    } else {
                        LOGGER.warn(new ParameterizedMessage("unexpected failure while sending request [{}] to [{}] for shard entry [{}]", actionName, masterNode, request), exp);
                        listener.onFailure(exp instanceof RemoteTransportException ? (Exception) (exp.getCause() instanceof Exception ? exp.getCause() : new ElasticsearchException(exp.getCause())) : exp);
                    }
                }
            });
    }
}