Examples with AsynchronousSearchContextEventListener org.opensearch.search.asynchronous.listener.AsynchronousSearchContextEventListener used on opensource projects

Example 1 with AsynchronousSearchContextEventListener

use of org.opensearch.search.asynchronous.listener.AsynchronousSearchContextEventListener in project asynchronous-search by opensearch-project.

the class AsynchronousSearchActiveStoreTests method testPutContextRejection.

public void testPutContextRejection() throws InterruptedException, BrokenBarrierException, TimeoutException {
    DiscoveryNode node = new DiscoveryNode("node", OpenSearchTestCase.buildNewFakeTransportAddress(), Collections.emptyMap(), DiscoveryNodeRole.BUILT_IN_ROLES, Version.CURRENT);
    ThreadPool testThreadPool = null;
    try {
        testThreadPool = new TestThreadPool(this.getClass().getName(), executorBuilder);
        AtomicInteger runningContexts = new AtomicInteger();
        ClusterService mockClusterService = ClusterServiceUtils.createClusterService(testThreadPool, node, clusterSettings);
        AsynchronousSearchActiveStore activeStore = new AsynchronousSearchActiveStore(mockClusterService);
        List<Runnable> runnables = new ArrayList<>();
        AtomicInteger numRejected = new AtomicInteger();
        AtomicInteger numFailures = new AtomicInteger();
        AtomicInteger numSuccesses = new AtomicInteger();
        int numThreads = maxRunningContexts + 1;
        CyclicBarrier barrier = new CyclicBarrier(numThreads + 1);
        CountDownLatch assertsLatch = new CountDownLatch(1);
        for (int i = 0; i < numThreads; i++) {
            ThreadPool finalTestThreadPool1 = testThreadPool;
            runnables.add(() -> {
                try {
                    AsynchronousSearchProgressListener asProgressListener = mockAsynchronousSearchProgressListener(finalTestThreadPool1);
                    AsynchronousSearchContextId asContextId = new AsynchronousSearchContextId(UUID.randomUUID().toString(), runningContexts.incrementAndGet());
                    boolean keepOnCompletion = randomBoolean();
                    TimeValue keepAlive = TimeValue.timeValueDays(randomInt(100));
                    AsynchronousSearchContextEventListener asContextEventListener = new AsynchronousSearchContextEventListener() {
                    };
                    AsynchronousSearchActiveContext context = new AsynchronousSearchActiveContext(asContextId, node.getId(), keepAlive, keepOnCompletion, finalTestThreadPool1, finalTestThreadPool1::absoluteTimeInMillis, asProgressListener, null, () -> true);
                    activeStore.putContext(asContextId, context, asContextEventListener::onContextRejected);
                    numSuccesses.getAndIncrement();
                    Optional<AsynchronousSearchActiveContext> optional = activeStore.getContext(asContextId);
                    assert (optional.isPresent());
                    assertEquals(optional.get(), context);
                    barrier.await(5, TimeUnit.SECONDS);
                    assertsLatch.await();
                    activeStore.freeContext(context.getContextId());
                    assertFalse(activeStore.getContext(context.getContextId()).isPresent());
                    barrier.await();
                } catch (OpenSearchRejectedExecutionException e) {
                    numRejected.getAndIncrement();
                    try {
                        barrier.await();
                        barrier.await();
                    } catch (InterruptedException | BrokenBarrierException ex) {
                        numFailures.getAndIncrement();
                    }
                } catch (InterruptedException | BrokenBarrierException e) {
                    numFailures.getAndIncrement();
                } catch (Exception e) {
                    logger.error(e.getMessage(), e);
                    numFailures.getAndIncrement();
                }
            });
        }
        ThreadPool finalTestThreadPool = testThreadPool;
        runnables.forEach(r -> finalTestThreadPool.generic().execute(r));
        // create contexts
        barrier.await(5, TimeUnit.SECONDS);
        assertEquals(activeStore.getAllContexts().size(), maxRunningContexts);
        assertEquals(numFailures.get(), 0);
        assertEquals(numRejected.get(), 1);
        assertEquals(numSuccesses.get(), maxRunningContexts);
        assertsLatch.countDown();
        // free contexts
        barrier.await(5, TimeUnit.SECONDS);
        assertEquals(activeStore.getAllContexts().size(), 0);
        assertEquals(numFailures.get(), 0);
    } finally {
        ThreadPool.terminate(testThreadPool, 10, TimeUnit.SECONDS);
    }
}

Example 2 with AsynchronousSearchContextEventListener

use of org.opensearch.search.asynchronous.listener.AsynchronousSearchContextEventListener in project asynchronous-search by opensearch-project.

the class AsynchronousSearchStateMachine method trigger.

/**
 * Triggers transition from current state on receiving an event. Also invokes {@linkplain Transition#onEvent()} and
 * {@linkplain Transition#eventListener()}.
 *
 * @param event to fire
 * @return The final Async search state
 * @throws AsynchronousSearchStateMachineClosedException the state machine has reached a terminal state
 */
@Override
public AsynchronousSearchState trigger(AsynchronousSearchContextEvent event) throws AsynchronousSearchStateMachineClosedException {
    AsynchronousSearchContext asynchronousSearchContext = event.asynchronousSearchContext();
    synchronized (asynchronousSearchContext) {
        AsynchronousSearchState currentState = asynchronousSearchContext.getAsynchronousSearchState();
        if (getFinalStates().contains(currentState)) {
            throw new AsynchronousSearchStateMachineClosedException(currentState, event);
        }
        String transitionId = getTransitionId(currentState, event.getClass());
        if (transitionsMap.containsKey(transitionId)) {
            AsynchronousSearchTransition<? extends AsynchronousSearchContextEvent> transition = transitionsMap.get(transitionId);
            execute(transition.onEvent(), event, currentState);
            asynchronousSearchContext.setState(transition.targetState());
            logger.debug("Executed event [{}] for asynchronous search id [{}] ", event.getClass().getName(), event.asynchronousSearchContext.getAsynchronousSearchId());
            BiConsumer<AsynchronousSearchContextId, AsynchronousSearchContextEventListener> eventListener = transition.eventListener();
            try {
                eventListener.accept(event.asynchronousSearchContext().getContextId(), asynchronousSearchContextEventListener);
            } catch (Exception ex) {
                logger.error(() -> new ParameterizedMessage("Failed to execute listener for asynchronous search id : [{}]", event.asynchronousSearchContext.getAsynchronousSearchId()), ex);
            }
            return asynchronousSearchContext.getAsynchronousSearchState();
        } else {
            String message = String.format(Locale.ROOT, "Invalid transition for " + "asynchronous search context [%s] from source state [%s] on event [%s]", asynchronousSearchContext.getAsynchronousSearchId(), currentState, event.getClass().getName());
            logger.error(message);
            throw new IllegalStateException(message);
        }
    }
}