Examples with TestCircuitBreaker org.opensearch.common.breaker.TestCircuitBreaker used on opensource projects

Example 1 with TestCircuitBreaker

use of org.opensearch.common.breaker.TestCircuitBreaker in project OpenSearch by opensearch-project.

the class InboundAggregatorTests method setUp.

@Before
@Override
public void setUp() throws Exception {
    super.setUp();
    Predicate<String> requestCanTripBreaker = action -> {
        if (unknownAction.equals(action)) {
            throw new ActionNotFoundTransportException(action);
        } else {
            return unBreakableAction.equals(action) == false;
        }
    };
    circuitBreaker = new TestCircuitBreaker();
    aggregator = new InboundAggregator(() -> circuitBreaker, requestCanTripBreaker);
}

Example 2 with TestCircuitBreaker

use of org.opensearch.common.breaker.TestCircuitBreaker in project OpenSearch by opensearch-project.

the class InboundPipelineTests method testPipelineHandling.

public void testPipelineHandling() throws IOException {
    final List<Tuple<MessageData, Exception>> expected = new ArrayList<>();
    final List<Tuple<MessageData, Exception>> actual = new ArrayList<>();
    final List<ReleasableBytesReference> toRelease = new ArrayList<>();
    final BiConsumer<TcpChannel, InboundMessage> messageHandler = (c, m) -> {
        try {
            final Header header = m.getHeader();
            final MessageData actualData;
            final Version version = header.getVersion();
            final boolean isRequest = header.isRequest();
            final long requestId = header.getRequestId();
            final boolean isCompressed = header.isCompressed();
            if (m.isShortCircuit()) {
                actualData = new MessageData(version, requestId, isRequest, isCompressed, header.getActionName(), null);
            } else if (isRequest) {
                final TestRequest request = new TestRequest(m.openOrGetStreamInput());
                actualData = new MessageData(version, requestId, isRequest, isCompressed, header.getActionName(), request.value);
            } else {
                final TestResponse response = new TestResponse(m.openOrGetStreamInput());
                actualData = new MessageData(version, requestId, isRequest, isCompressed, null, response.value);
            }
            actual.add(new Tuple<>(actualData, m.getException()));
        } catch (IOException e) {
            throw new AssertionError(e);
        }
    };
    final StatsTracker statsTracker = new StatsTracker();
    final LongSupplier millisSupplier = () -> TimeValue.nsecToMSec(System.nanoTime());
    final InboundDecoder decoder = new InboundDecoder(Version.CURRENT, PageCacheRecycler.NON_RECYCLING_INSTANCE);
    final String breakThisAction = "break_this_action";
    final String actionName = "actionName";
    final Predicate<String> canTripBreaker = breakThisAction::equals;
    final TestCircuitBreaker circuitBreaker = new TestCircuitBreaker();
    circuitBreaker.startBreaking();
    final InboundAggregator aggregator = new InboundAggregator(() -> circuitBreaker, canTripBreaker);
    final InboundPipeline pipeline = new InboundPipeline(statsTracker, millisSupplier, decoder, aggregator, messageHandler);
    final FakeTcpChannel channel = new FakeTcpChannel();
    final int iterations = randomIntBetween(5, 10);
    long totalMessages = 0;
    long bytesReceived = 0;
    for (int i = 0; i < iterations; ++i) {
        actual.clear();
        expected.clear();
        toRelease.clear();
        try (BytesStreamOutput streamOutput = new BytesStreamOutput()) {
            while (streamOutput.size() < BYTE_THRESHOLD) {
                final Version version = randomFrom(Version.CURRENT, Version.CURRENT.minimumCompatibilityVersion());
                final String value = randomAlphaOfLength(randomIntBetween(10, 200));
                final boolean isRequest = randomBoolean();
                final boolean isCompressed = randomBoolean();
                final long requestId = totalMessages++;
                final MessageData messageData;
                Exception expectedExceptionClass = null;
                OutboundMessage message;
                if (isRequest) {
                    if (rarely()) {
                        messageData = new MessageData(version, requestId, true, isCompressed, breakThisAction, null);
                        message = new OutboundMessage.Request(threadContext, new String[0], new TestRequest(value), version, breakThisAction, requestId, false, isCompressed);
                        expectedExceptionClass = new CircuitBreakingException("", CircuitBreaker.Durability.PERMANENT);
                    } else {
                        messageData = new MessageData(version, requestId, true, isCompressed, actionName, value);
                        message = new OutboundMessage.Request(threadContext, new String[0], new TestRequest(value), version, actionName, requestId, false, isCompressed);
                    }
                } else {
                    messageData = new MessageData(version, requestId, false, isCompressed, null, value);
                    message = new OutboundMessage.Response(threadContext, Collections.emptySet(), new TestResponse(value), version, requestId, false, isCompressed);
                }
                expected.add(new Tuple<>(messageData, expectedExceptionClass));
                final BytesReference reference = message.serialize(new BytesStreamOutput());
                Streams.copy(reference.streamInput(), streamOutput);
            }
            final BytesReference networkBytes = streamOutput.bytes();
            int currentOffset = 0;
            while (currentOffset != networkBytes.length()) {
                final int remainingBytes = networkBytes.length() - currentOffset;
                final int bytesToRead = Math.min(randomIntBetween(1, 32 * 1024), remainingBytes);
                final BytesReference slice = networkBytes.slice(currentOffset, bytesToRead);
                try (ReleasableBytesReference reference = new ReleasableBytesReference(slice, () -> {
                })) {
                    toRelease.add(reference);
                    bytesReceived += reference.length();
                    pipeline.handleBytes(channel, reference);
                    currentOffset += bytesToRead;
                }
            }
            final int messages = expected.size();
            for (int j = 0; j < messages; ++j) {
                final Tuple<MessageData, Exception> expectedTuple = expected.get(j);
                final Tuple<MessageData, Exception> actualTuple = actual.get(j);
                final MessageData expectedMessageData = expectedTuple.v1();
                final MessageData actualMessageData = actualTuple.v1();
                assertEquals(expectedMessageData.requestId, actualMessageData.requestId);
                assertEquals(expectedMessageData.isRequest, actualMessageData.isRequest);
                assertEquals(expectedMessageData.isCompressed, actualMessageData.isCompressed);
                assertEquals(expectedMessageData.actionName, actualMessageData.actionName);
                assertEquals(expectedMessageData.value, actualMessageData.value);
                if (expectedTuple.v2() != null) {
                    assertNotNull(actualTuple.v2());
                    assertThat(actualTuple.v2(), instanceOf(expectedTuple.v2().getClass()));
                }
            }
            for (ReleasableBytesReference released : toRelease) {
                assertEquals(0, released.refCount());
            }
        }
        assertEquals(bytesReceived, statsTracker.getBytesRead());
        assertEquals(totalMessages, statsTracker.getMessagesReceived());
    }
}