SizedBufferingTraceConsumer implements FlushableTraceConsumer { @Override public void receive(Traces traces) { synchronized(monitor) { for (Trace trace : traces.getTracesList()) { size += traceSizer.size(trace); traceBuffer.put(trace); if (size >= bufferSize) { flush(); } } } } SizedBufferingTraceConsumer(TraceConsumer traceConsumer, Sizer<Trace> traceSizer, int bufferSize); @Override void receive(Traces traces); @Override void flush();  }
@Test public void testBuffering() { ArgumentCaptor<Traces> tracesCaptor = ArgumentCaptor.forClass(Traces.class); TraceConsumer mockSink = mock(TraceConsumer.class); final Trace trace1 = Trace.newBuilder().setProjectId("1").setTraceId("1").build(); final Trace trace2 = Trace.newBuilder().setProjectId("1").setTraceId("2").build(); final Trace trace3 = Trace.newBuilder().setProjectId("1").setTraceId("3").build(); Sizer<Trace> testSizer = new Sizer<Trace>() { @Override public int size(Trace sizeable) { if (sizeable == trace1) { return 10; } else if (sizeable == trace2) { return 20; } else if (sizeable == trace3) { return 30; } else { return 0; } } }; FlushableTraceConsumer bufferedSink = new SizedBufferingTraceConsumer(mockSink, testSizer, 25); bufferedSink.receive(Traces.newBuilder().addTraces(trace1).build()); bufferedSink.receive(Traces.newBuilder().addTraces(trace2).build()); bufferedSink.receive(Traces.newBuilder().addTraces(trace3).build()); verify(mockSink, times(2)).receive(tracesCaptor.capture()); List<Traces> traceBatches = tracesCaptor.getAllValues(); assertThat(traceBatches).hasSize(2); assertThat(traceBatches.get(0).getTracesList()).containsExactly(trace1, trace2); assertThat(traceBatches.get(1).getTracesList()).containsExactly(trace3); }