use of io.pravega.shared.protocol.netty.WireCommands.ConditionalAppend in project pravega by pravega.
the class ConditionalOutputStreamImpl method write.
@Override
public boolean write(ByteBuffer data, long expectedOffset) throws SegmentSealedException {
Exceptions.checkNotClosed(closed.get(), this);
synchronized (lock) {
// Used to preserver order.
long appendSequence = requestIdGenerator.get();
return retrySchedule.retryWhen(e -> {
Throwable cause = Exceptions.unwrap(e);
boolean hasTokenExpired = cause instanceof TokenExpiredException;
if (hasTokenExpired) {
this.tokenProvider.signalTokenExpired();
}
return cause instanceof Exception && (hasTokenExpired || cause instanceof ConnectionFailedException);
}).run(() -> {
if (client == null || client.isClosed()) {
client = new RawClient(controller, connectionPool, segmentId);
long requestId = client.getFlow().getNextSequenceNumber();
log.debug("Setting up appends on segment {} for ConditionalOutputStream with writer id {}", segmentId, writerId);
CompletableFuture<Reply> reply = tokenProvider.retrieveToken().thenCompose(token -> {
SetupAppend setup = new SetupAppend(requestId, writerId, segmentId.getScopedName(), token);
return client.sendRequest(requestId, setup);
});
AppendSetup appendSetup = transformAppendSetup(reply.join());
if (appendSetup.getLastEventNumber() >= appendSequence) {
return true;
}
}
long requestId = client.getFlow().getNextSequenceNumber();
val request = new ConditionalAppend(writerId, appendSequence, expectedOffset, new Event(Unpooled.wrappedBuffer(data)), requestId);
val reply = client.sendRequest(requestId, request);
return transformDataAppended(reply.join());
});
}
}
use of io.pravega.shared.protocol.netty.WireCommands.ConditionalAppend in project pravega by pravega.
the class ConditionalOutputStreamTest method testRetriesOnTokenExpiry.
@SneakyThrows
@Test(timeout = 10000)
public void testRetriesOnTokenExpiry() {
@Cleanup MockConnectionFactoryImpl connectionFactory = new MockConnectionFactoryImpl();
@Cleanup MockController controller = new MockController("localhost", 0, connectionFactory, true);
ConditionalOutputStreamFactory factory = new ConditionalOutputStreamFactoryImpl(controller, connectionFactory);
Segment segment = new Segment("scope", "testWrite", 1);
@Cleanup ConditionalOutputStream objectUnderTest = factory.createConditionalOutputStream(segment, DelegationTokenProviderFactory.create("token", controller, segment, AccessOperation.ANY), EventWriterConfig.builder().build());
ByteBuffer data = ByteBuffer.allocate(10);
ClientConnection clientConnection = Mockito.mock(ClientConnection.class);
PravegaNodeUri location = new PravegaNodeUri("localhost", 0);
connectionFactory.provideConnection(location, clientConnection);
setupAppend(connectionFactory, segment, clientConnection, location);
final AtomicLong retryCounter = new AtomicLong(0);
Mockito.doAnswer(new Answer<Void>() {
@Override
public Void answer(InvocationOnMock invocation) throws Throwable {
ConditionalAppend argument = (ConditionalAppend) invocation.getArgument(0);
ReplyProcessor processor = connectionFactory.getProcessor(location);
if (retryCounter.getAndIncrement() < 2) {
processor.process(new WireCommands.AuthTokenCheckFailed(argument.getRequestId(), "SomeException", WireCommands.AuthTokenCheckFailed.ErrorCode.TOKEN_EXPIRED));
} else {
processor.process(new WireCommands.DataAppended(argument.getRequestId(), argument.getWriterId(), argument.getEventNumber(), 0, -1));
}
return null;
}
}).when(clientConnection).send(any(ConditionalAppend.class));
assertTrue(objectUnderTest.write(data, 0));
assertEquals(3, retryCounter.get());
}
use of io.pravega.shared.protocol.netty.WireCommands.ConditionalAppend in project pravega by pravega.
the class ConditionalOutputStreamTest method testWrite.
@Test(timeout = 10000)
public void testWrite() throws ConnectionFailedException, SegmentSealedException {
@Cleanup MockConnectionFactoryImpl connectionFactory = new MockConnectionFactoryImpl();
@Cleanup MockController controller = new MockController("localhost", 0, connectionFactory, true);
ConditionalOutputStreamFactory factory = new ConditionalOutputStreamFactoryImpl(controller, connectionFactory);
Segment segment = new Segment("scope", "testWrite", 1);
@Cleanup ConditionalOutputStream cOut = factory.createConditionalOutputStream(segment, DelegationTokenProviderFactory.create("token", controller, segment, AccessOperation.ANY), EventWriterConfig.builder().build());
ByteBuffer data = ByteBuffer.allocate(10);
ClientConnection mock = Mockito.mock(ClientConnection.class);
PravegaNodeUri location = new PravegaNodeUri("localhost", 0);
connectionFactory.provideConnection(location, mock);
setupAppend(connectionFactory, segment, mock, location);
Mockito.doAnswer(new Answer<Void>() {
@Override
public Void answer(InvocationOnMock invocation) throws Throwable {
ConditionalAppend argument = (ConditionalAppend) invocation.getArgument(0);
ReplyProcessor processor = connectionFactory.getProcessor(location);
if (argument.getExpectedOffset() == 0 || argument.getExpectedOffset() == 2) {
processor.process(new WireCommands.DataAppended(argument.getRequestId(), argument.getWriterId(), argument.getEventNumber(), 0, -1));
} else {
processor.process(new WireCommands.ConditionalCheckFailed(argument.getWriterId(), argument.getEventNumber(), argument.getRequestId()));
}
return null;
}
}).when(mock).send(any(ConditionalAppend.class));
assertTrue(cOut.write(data, 0));
assertFalse(cOut.write(data, 1));
assertTrue(cOut.write(data, 2));
assertFalse(cOut.write(data, 3));
}
use of io.pravega.shared.protocol.netty.WireCommands.ConditionalAppend in project pravega by pravega.
the class ConditionalOutputStreamTest method testNonExpiryTokenCheckFailure.
@Test(timeout = 10000)
public void testNonExpiryTokenCheckFailure() throws ConnectionFailedException {
@Cleanup MockConnectionFactoryImpl connectionFactory = new MockConnectionFactoryImpl();
@Cleanup MockController controller = new MockController("localhost", 0, connectionFactory, true);
ConditionalOutputStreamFactory factory = new ConditionalOutputStreamFactoryImpl(controller, connectionFactory);
Segment segment = new Segment("scope", "testWrite", 1);
@Cleanup ConditionalOutputStream objectUnderTest = factory.createConditionalOutputStream(segment, DelegationTokenProviderFactory.create("token", controller, segment, AccessOperation.ANY), EventWriterConfig.builder().build());
ByteBuffer data = ByteBuffer.allocate(10);
ClientConnection clientConnection = Mockito.mock(ClientConnection.class);
PravegaNodeUri location = new PravegaNodeUri("localhost", 0);
connectionFactory.provideConnection(location, clientConnection);
setupAppend(connectionFactory, segment, clientConnection, location);
Mockito.doAnswer(new Answer<Void>() {
@Override
public Void answer(InvocationOnMock invocation) throws Throwable {
ConditionalAppend argument = (ConditionalAppend) invocation.getArgument(0);
ReplyProcessor processor = connectionFactory.getProcessor(location);
processor.process(new WireCommands.AuthTokenCheckFailed(argument.getRequestId(), "SomeException", WireCommands.AuthTokenCheckFailed.ErrorCode.TOKEN_CHECK_FAILED));
return null;
}
}).when(clientConnection).send(any(ConditionalAppend.class));
AssertExtensions.assertThrows(AuthenticationException.class, () -> objectUnderTest.write(data, 0));
}
use of io.pravega.shared.protocol.netty.WireCommands.ConditionalAppend in project pravega by pravega.
the class RawClientTest method testRequestReply.
@Test
public void testRequestReply() throws ConnectionFailedException, InterruptedException, ExecutionException {
PravegaNodeUri endpoint = new PravegaNodeUri("localhost", -1);
@Cleanup MockConnectionFactoryImpl connectionFactory = new MockConnectionFactoryImpl();
@Cleanup MockController controller = new MockController(endpoint.getEndpoint(), endpoint.getPort(), connectionFactory);
ClientConnection connection = Mockito.mock(ClientConnection.class);
connectionFactory.provideConnection(endpoint, connection);
@Cleanup RawClient rawClient = new RawClient(controller, connectionFactory, new Segment("scope", "testHello", 0));
UUID id = UUID.randomUUID();
ConditionalAppend request = new ConditionalAppend(id, 1, 0, Unpooled.EMPTY_BUFFER);
CompletableFuture<Reply> future = rawClient.sendRequest(1, request);
Mockito.verify(connection).send(request);
assertFalse(future.isDone());
ReplyProcessor processor = connectionFactory.getProcessor(endpoint);
DataAppended reply = new DataAppended(id, 1, 0);
processor.process(reply);
assertTrue(future.isDone());
assertEquals(reply, future.get());
}
Aggregations