/** * Write a whole batch to eventhub */ public Future<WriteResponse> write (Batch<String> batch, WriteCallback callback) { Timer.Context context = writeTimer.time(); int returnCode = 0; LOG.info ("Dispatching batch " + batch.getId()); recordsAttempted.mark(batch.getRecords().size()); try { String encoded = encodeBatch(batch); returnCode = request (encoded); WriteResponse<Integer> response = WRITE_RESPONSE_WRAPPER.wrap(returnCode); callback.onSuccess(response); bytesWritten.mark(encoded.length()); recordsSuccess.mark(batch.getRecords().size()); } catch (Exception e) { LOG.error("Dispatching batch " + batch.getId() + " failed :" + e.toString()); callback.onFailure(e); recordsFailed.mark(batch.getRecords().size()); } context.close(); Future<Integer> future = Futures.immediateFuture(returnCode); return new WriteResponseFuture<>(future, WRITE_RESPONSE_WRAPPER); }
if (!hasRoom(record, largeMessagePolicy)) { LOG.debug ("Cannot add {} to previous batch because the batch already has {} bytes", record.toString(), getCurrentSizeInByte()); if (largeMessagePolicy == LargeMessagePolicy.FAIL) { throw new RecordTooLargeException(); this.append(record); thunks.add(new Thunk(callback, getRecordSizeInByte(record))); RecordFuture future = new RecordFuture(latch, recordCount); recordCount++;
/** * This will block until all the incomplete batches are acknowledged */ public void flush() { try { ArrayList<Batch> batches = this.incomplete.all(); int numOutstandingRecords = 0; for (Batch batch: batches) { numOutstandingRecords += batch.getRecords().size(); } LOG.debug ("Flush called on {} batches with {} records total", batches.size(), numOutstandingRecords); for (Batch batch: batches) { batch.await(); } } catch (Exception e) { LOG.error ("Error happened while flushing batches"); } }
/** * Each record of batch is wrapped by a 'Body' json object * put this new object into an array, encode the whole array */ private String encodeBatch (Batch<String> batch) throws IOException { // Convert original json object to a new json object with format {"Body": "originalJson"} // Add new json object to an array and send the whole array to eventhub using REST api // Refer to https://docs.microsoft.com/en-us/rest/api/eventhub/send-batch-events List<String> records = batch.getRecords(); ArrayList<EventhubRequest> arrayList = new ArrayList<>(); for (String record: records) { arrayList.add(new EventhubRequest(record)); } return mapper.writeValueAsString (arrayList); }
@Test public void testSingleBatch() { // mock eventhub data writer Properties props = new Properties(); EventhubDataWriter eventhubDataWriter = Mockito.spy(new EventhubDataWriter(props, mockHttpClient)); Mockito.doNothing().when(eventhubDataWriter).refreshSignature(); List<String> records = new LinkedList<>(); for (int i=0; i<50; ++i) records.add(new String("abcdefgh")); Batch<String> batch = mock(Batch.class); WriteCallback callback = mock(WriteCallback.class); Mockito.when(batch.getRecords()).thenReturn(records); Future<WriteResponse> future = eventhubDataWriter.write(batch,callback); verify(callback, times(1)).onSuccess(isA(WriteResponse.class)); verify(callback, never()).onFailure(isA(Exception.class)); Assert.assertTrue(future.isDone(), "Future should be done"); }
protected Pair<BulkRequest, FutureCallbackHolder> prepareBatch(Batch<Object> batch, WriteCallback callback) { BulkRequest bulkRequest = new BulkRequest(); final StringBuilder stringBuilder = new StringBuilder(); for (Object record : batch.getRecords()) { try { byte[] serializedBytes = this.serializer.serializeToJson(record); log.debug("serialized record: {}", serializedBytes); IndexRequest indexRequest = new IndexRequest(this.indexName, this.indexType) .source(serializedBytes, 0, serializedBytes.length, XContentType.JSON); if (this.idMappingEnabled) { String id = this.typeMapper.getValue(this.idFieldName, record); indexRequest.id(id); stringBuilder.append(";").append(id); } bulkRequest.add(indexRequest); } catch (Exception e) { log.error("Encountered exception {}", e); } } FutureCallbackHolder futureCallbackHolder = new FutureCallbackHolder(callback, exception -> log.error("Batch: {} failed on ids; {} with exception {}", batch.getId(), stringBuilder.toString(), exception), this.malformedDocPolicy); return new Pair(bulkRequest, futureCallbackHolder); }
@Test public void testClose () throws IOException, InterruptedException { SequentialBasedBatchAccumulator accumulator = new EventhubBatchAccumulator(64, 3000, 5); StringBuffer buffer = new StringBuffer(); for (int i = 0; i < 40; ++i) { buffer.append('a'); } String record1 = buffer.toString() + "1"; String record2 = buffer.toString() + "2"; String record3 = buffer.toString() + "3"; String record4 = buffer.toString() + "4"; String record5 = buffer.toString() + "5"; accumulator.append(record1, WriteCallback.EMPTY); accumulator.append(record2, WriteCallback.EMPTY); accumulator.append(record3, WriteCallback.EMPTY); accumulator.append(record4, WriteCallback.EMPTY); accumulator.append(record5, WriteCallback.EMPTY); (new Thread(new CloseAccumulatorThread(accumulator))).start(); Thread.sleep(1000); Assert.assertEquals(accumulator.getNextAvailableBatch().getRecords().get(0), record1); Assert.assertEquals(accumulator.getNextAvailableBatch().getRecords().get(0), record2); Assert.assertEquals(accumulator.getNextAvailableBatch().getRecords().get(0), record3); Assert.assertEquals(accumulator.getNextAvailableBatch().getRecords().get(0), record4); Assert.assertEquals(accumulator.getNextAvailableBatch().getRecords().get(0), record5); }
/** * This will block until all the incomplete batches are acknowledged */ public void flush() { try { ArrayList<Batch> batches = this.incomplete.all(); int numOutstandingRecords = 0; for (Batch batch: batches) { numOutstandingRecords += batch.getRecords().size(); } LOG.debug ("Flush called on {} batches with {} records total", batches.size(), numOutstandingRecords); for (Batch batch: batches) { batch.await(); } } catch (Exception e) { LOG.error ("Error happened while flushing batches"); } }
if (!hasRoom(record, largeMessagePolicy)) { LOG.debug ("Cannot add {} to previous batch because the batch already has {} bytes", record.toString(), getCurrentSizeInByte()); if (largeMessagePolicy == LargeMessagePolicy.FAIL) { throw new RecordTooLargeException(); this.append(record); thunks.add(new Thunk(callback, getRecordSizeInByte(record))); RecordFuture future = new RecordFuture(latch, recordCount); recordCount++;
/** * Write a whole batch to eventhub */ public Future<WriteResponse> write (Batch<String> batch, WriteCallback callback) { Timer.Context context = writeTimer.time(); int returnCode = 0; LOG.info ("Dispatching batch " + batch.getId()); recordsAttempted.mark(batch.getRecords().size()); try { String encoded = encodeBatch(batch); returnCode = request (encoded); WriteResponse<Integer> response = WRITE_RESPONSE_WRAPPER.wrap(returnCode); callback.onSuccess(response); bytesWritten.mark(encoded.length()); recordsSuccess.mark(batch.getRecords().size()); } catch (Exception e) { LOG.error("Dispatching batch " + batch.getId() + " failed :" + e.toString()); callback.onFailure(e); recordsFailed.mark(batch.getRecords().size()); } context.close(); Future<Integer> future = Futures.immediateFuture(returnCode); return new WriteResponseFuture<>(future, WRITE_RESPONSE_WRAPPER); }
/** * Each record of batch is wrapped by a 'Body' json object * put this new object into an array, encode the whole array */ private String encodeBatch (Batch<String> batch) throws IOException { // Convert original json object to a new json object with format {"Body": "originalJson"} // Add new json object to an array and send the whole array to eventhub using REST api // Refer to https://docs.microsoft.com/en-us/rest/api/eventhub/send-batch-events List<String> records = batch.getRecords(); ArrayList<EventhubRequest> arrayList = new ArrayList<>(); for (String record: records) { arrayList.add(new EventhubRequest(record)); } return mapper.writeValueAsString (arrayList); }
protected Pair<BulkRequest, FutureCallbackHolder> prepareBatch(Batch<Object> batch, WriteCallback callback) { BulkRequest bulkRequest = new BulkRequest(); final StringBuilder stringBuilder = new StringBuilder(); for (Object record : batch.getRecords()) { try { byte[] serializedBytes = this.serializer.serializeToJson(record); log.debug("serialized record: {}", serializedBytes); IndexRequest indexRequest = new IndexRequest(this.indexName, this.indexType) .source(serializedBytes, 0, serializedBytes.length, XContentType.JSON); if (this.idMappingEnabled) { String id = this.typeMapper.getValue(this.idFieldName, record); indexRequest.id(id); stringBuilder.append(";").append(id); } bulkRequest.add(indexRequest); } catch (Exception e) { log.error("Encountered exception {}", e); } } FutureCallbackHolder futureCallbackHolder = new FutureCallbackHolder(callback, exception -> log.error("Batch: {} failed on ids; {} with exception {}", batch.getId(), stringBuilder.toString(), exception), this.malformedDocPolicy); return new Pair(bulkRequest, futureCallbackHolder); }