Examples with TransportBuffer com.nextdoor.bender.ipc.TransportBuffer used on opensource projects

Example 1 with TransportBuffer

use of com.nextdoor.bender.ipc.TransportBuffer in project bender by Nextdoor.

the class BaseHandlerTest method testIpcOnAddFailure.

@Test
public void testIpcOnAddFailure() throws Throwable {
    BaseHandler.CONFIG_FILE = "/config/handler_config.json";
    handler.skipWriteStats = true;
    List<DummyEvent> events = new ArrayList<DummyEvent>(2);
    events.add(new DummyEvent("foo", 0));
    events.add(new DummyEvent("bar", 0));
    TestContext context = new TestContext();
    context.setInvokedFunctionArn("arn:aws:lambda:us-east-1:123:function:test:tag");
    handler.init(context);
    TransportBuffer tbSpy1 = spy(new ArrayTransportBuffer());
    TransportBuffer tbSpy2 = spy(new ArrayTransportBuffer());
    doCallRealMethod().doCallRealMethod().when(tbSpy1).add(any());
    doThrow(new IllegalStateException("expected")).when(tbSpy2).add(any());
    IpcSenderService spyIpc = spy(handler.getIpcService());
    TransportFactory tfSpy = spy(spyIpc.getTransportFactory());
    when(tfSpy.newTransportBuffer()).thenReturn(tbSpy1, tbSpy2);
    spyIpc.setTransportFactory(tfSpy);
    handler.setIpcService(spyIpc);
    handler.handler(events, context);
    assertEquals(1, spyIpc.getSuccessCountStat().getValue());
}

Example 2 with TransportBuffer

use of com.nextdoor.bender.ipc.TransportBuffer in project bender by Nextdoor.

the class S3Transport method sendBatch.

@Override
public void sendBatch(TransportBuffer buffer, LinkedHashMap<String, String> partitions, Context context) throws TransportException {
    S3TransportBuffer buf = (S3TransportBuffer) buffer;
    /*
     * Create s3 key (filepath + filename)
     */
    LinkedHashMap<String, String> parts = new LinkedHashMap<String, String>(partitions);
    String filename = parts.remove(FILENAME_KEY);
    if (filename == null) {
        filename = context.getAwsRequestId();
    }
    String key = parts.entrySet().stream().map(s -> s.getKey() + "=" + s.getValue()).collect(Collectors.joining("/"));
    key = (key.equals("") ? filename : key + '/' + filename);
    if (this.basePath.endsWith("/")) {
        key = this.basePath + key;
    } else {
        key = this.basePath + '/' + key;
    }
    // TODO: make this dynamic
    if (key.endsWith(".gz")) {
        key = key.substring(0, key.length() - 3);
    }
    /*
     * Add or strip out compression format extension
     *
     * TODO: get this based on the compression codec
     */
    if (this.compress || buf.isCompressed()) {
        key += ".bz2";
    }
    ByteArrayOutputStream os = buf.getInternalBuffer();
    /*
     * Compress stream if needed. Don't compress a compressed stream.
     */
    ByteArrayOutputStream payload;
    if (this.compress && !buf.isCompressed()) {
        payload = compress(os);
    } else {
        payload = os;
    }
    /*
     * For memory efficiency convert the output stream into an InputStream. This is done using the
     * easystream library but under the hood it uses piped streams to facilitate this process. This
     * avoids copying the entire contents of the OutputStream to populate the InputStream. Note that
     * this process creates another thread to consume from the InputStream.
     */
    final String s3Key = key;
    /*
     * Write to OutputStream
     */
    final InputStreamFromOutputStream<String> isos = new InputStreamFromOutputStream<String>() {

        public String produce(final OutputStream dataSink) throws Exception {
            /*
         * Note this is executed in a different thread
         */
            payload.writeTo(dataSink);
            return null;
        }
    };
    /*
     * Consume InputStream
     */
    try {
        sendStream(isos, s3Key, payload.size());
    } finally {
        try {
            isos.close();
        } catch (IOException e) {
            throw new TransportException(e);
        } finally {
            buf.close();
        }
    }
}

Example 3 with TransportBuffer

use of com.nextdoor.bender.ipc.TransportBuffer in project bender by Nextdoor.

the class TcpTransport method sendBatch.

@Override
public void sendBatch(TransportBuffer buffer) throws TransportException {
    Buffer internalBuffer = ((TcpTransportBuffer) buffer).getInternalBuffer();
    Callable<Boolean> write = () -> {
        try {
            sink.write(internalBuffer, internalBuffer.size());
            return true;
        } catch (IOException ex) {
            throw new TransportException("Error while sending in tcp transport", ex);
        }
    };
    try {
        new CallExecutor(retryConfig).execute(write);
    } catch (RetriesExhaustedException | UnexpectedException ue) {
        throw new TransportException(ue);
    }
}