Examples with BagheeraMessage com.mozilla.bagheera.BagheeraProto.BagheeraMessage used on opensource projects

Example 1 with BagheeraMessage

use of com.mozilla.bagheera.BagheeraProto.BagheeraMessage in project bagheera by mozilla-metrics.

the class KafkaConsumer method poll.

@Override
public void poll() {
    final CountDownLatch latch = new CountDownLatch(streams.size());
    for (final KafkaStream<Message> stream : streams) {
        workers.add(executor.submit(new Callable<Void>() {

            @Override
            public Void call() {
                try {
                    for (MessageAndMetadata<Message> mam : stream) {
                        BagheeraMessage bmsg = BagheeraMessage.parseFrom(ByteString.copyFrom(mam.message().payload()));
                        // get the sink for this message's namespace
                        // (typically only one sink unless a regex pattern was used to listen to multiple topics)
                        KeyValueSink sink = sinkFactory.getSink(bmsg.getNamespace());
                        if (sink == null) {
                            LOG.error("Could not obtain sink for namespace: " + bmsg.getNamespace());
                            break;
                        }
                        if (bmsg.getOperation() == Operation.CREATE_UPDATE && bmsg.hasId() && bmsg.hasPayload()) {
                            if (validationPipeline == null || validationPipeline.isValid(bmsg.getPayload().toByteArray())) {
                                if (bmsg.hasTimestamp()) {
                                    sink.store(bmsg.getId(), bmsg.getPayload().toByteArray(), bmsg.getTimestamp());
                                } else {
                                    sink.store(bmsg.getId(), bmsg.getPayload().toByteArray());
                                }
                            } else {
                                invalidMessageMeter.mark();
                                // TODO: sample out an example payload
                                LOG.warn("Invalid payload for namespace: " + bmsg.getNamespace());
                            }
                        } else if (bmsg.getOperation() == Operation.DELETE && bmsg.hasId()) {
                            sink.delete(bmsg.getId());
                        }
                        consumed.mark();
                    }
                } catch (InvalidProtocolBufferException e) {
                    LOG.error("Invalid protocol buffer in data stream", e);
                } catch (UnsupportedEncodingException e) {
                    LOG.error("Message ID was not in UTF-8 encoding", e);
                } catch (IOException e) {
                    LOG.error("IO error while storing to data sink", e);
                } finally {
                    latch.countDown();
                }
                return null;
            }
        }));
    }
    // run indefinitely unless we detect that a thread exited
    try {
        while (true) {
            latch.await(10, TimeUnit.SECONDS);
            if (latch.getCount() != streams.size()) {
                // we have a dead thread and should exit
                break;
            }
        }
    } catch (InterruptedException e) {
        LOG.info("Interrupted during polling", e);
    }
    // Spit out errors if there were any
    for (Future<Void> worker : workers) {
        try {
            if (worker.isDone() && !worker.isCancelled()) {
                worker.get(1, TimeUnit.SECONDS);
            }
        } catch (InterruptedException e) {
            LOG.error("Thread was interrupted:", e);
        } catch (ExecutionException e) {
            LOG.error("Exception occured in thread:", e);
        } catch (TimeoutException e) {
            LOG.error("Timed out waiting for thread result:", e);
        } catch (CancellationException e) {
            LOG.error("Thread has been canceled: ", e);
        }
    }
}

Example 2 with BagheeraMessage

use of com.mozilla.bagheera.BagheeraProto.BagheeraMessage in project bagheera by mozilla-metrics.

the class ProducerTest method produceData.

private void produceData(boolean includeBadRecord) throws InterruptedException {
    Properties props = getProperties();
    kafka.javaapi.producer.Producer<String, BagheeraMessage> producer = new kafka.javaapi.producer.Producer<String, BagheeraMessage>(new ProducerConfig(props));
    BagheeraMessage msg = getMessage(GOOD_MESSAGE_SIZE);
    assertEquals(GOOD_MESSAGE_SIZE, msg.getPayload().size());
    producer.send(getProducerData(msg));
    producer.send(getProducerData(getMessage(GOOD_MESSAGE_SIZE)));
    if (includeBadRecord) {
        producer.send(getProducerData(getMessage(BAD_MESSAGE_SIZE)));
    }
    for (int i = 0; i < BATCH_SIZE; i++) {
        producer.send(getProducerData(getMessage(GOOD_MESSAGE_SIZE)));
    }
    producer.close();
    // Wait for flush
    Thread.sleep(100);
}

Example 3 with BagheeraMessage

use of com.mozilla.bagheera.BagheeraProto.BagheeraMessage in project bagheera by mozilla-metrics.

the class SubmissionHandlerTest method testSetFields.

@Test
public void testSetFields() throws Exception {
    SubmissionHandler handler = new SubmissionHandler(null, null, null, null);
    BagheeraMessage.Builder builder = BagheeraMessage.newBuilder();
    BagheeraMessage before = builder.buildPartial();
    assertEquals("", before.getId());
    assertEquals("", before.getNamespace());
    assertEquals("", before.getApiVersion());
    assertEquals(0, before.getPartitionCount());
    assertEquals(0l, before.getTimestamp());
    String expectedNamespace = "test";
    String expectedApiVersion = "2.5";
    String expectedId = "hello there";
    long expectedTimestamp = System.currentTimeMillis();
    List<String> expectedPartitions = new ArrayList<String>();
    BagheeraHttpRequest request = Mockito.mock(BagheeraHttpRequest.class);
    Mockito.when(request.getNamespace()).thenReturn(expectedNamespace);
    Mockito.when(request.getApiVersion()).thenReturn(expectedApiVersion);
    Mockito.when(request.getId()).thenReturn(expectedId);
    Mockito.when(request.getPartitions()).thenReturn(expectedPartitions);
    // Make sure we don't interrogate the mocked InetSocketAddress below.
    Mockito.when(request.getHeader(HttpUtil.X_FORWARDED_FOR)).thenReturn("123.123.123.123");
    MessageEvent event = Mockito.mock(MessageEvent.class);
    Channel channel = Mockito.mock(Channel.class);
    Mockito.when(event.getChannel()).thenReturn(channel);
    InetSocketAddress address = Mockito.mock(InetSocketAddress.class);
    Mockito.when(channel.getRemoteAddress()).thenReturn(address);
    // Do not set the ID
    // handler.setMessageFields(request, event, builder, expectedTimestamp, false);
    BagheeraMessage after = builder.build();
    // <-- missing ID
    assertEquals("", after.getId());
    // assertEquals(expectedNamespace, after.getNamespace());
    // assertEquals(expectedApiVersion, after.getApiVersion());
    // assertEquals(0, after.getPartitionCount());
    // assertEquals(expectedTimestamp, after.getTimestamp());
    builder = BagheeraMessage.newBuilder();
// This time, *do* set the ID
// handler.setMessageFields(request, event, builder, expectedTimestamp, true);
// after = builder.build();
// assertEquals(expectedId, after.getId()); // <-- ID has been set.
// assertEquals(expectedNamespace, after.getNamespace());
// assertEquals(expectedApiVersion, after.getApiVersion());
// assertEquals(0, after.getPartitionCount());
// assertEquals(expectedTimestamp, after.getTimestamp());
// // Test without specifying an apiVersion
// Mockito.when(request.getApiVersion()).thenReturn(null);
// builder = BagheeraMessage.newBuilder();
// handler.setMessageFields(request, event, builder, expectedTimestamp, true);
// after = builder.build();
// assertEquals(expectedId, after.getId()); // <-- ID has been set.
// assertEquals(expectedNamespace, after.getNamespace());
// assertEquals("", after.getApiVersion());
// assertEquals(0, after.getPartitionCount());
// assertEquals(expectedTimestamp, after.getTimestamp());
// Test with some partitions
// Expectedpartitions.add("hello");
// expectedPartitions.add("goodbye");
// Mockito.when(request.getPartitions()).thenReturn(expectedPartitions);
// builder = BagheeraMessage.newBuilder();
// assertEquals(0, builder.getPartitionCount());
// handler.setMessageFields(request, event, builder, expectedTimestamp, true);
// after = builder.build();
// assertEquals(expectedPartitions.size(), after.getPartitionCount());
// for (int i = 0; i < expectedPartitions.size(); i++) {
// assertEquals(expectedPartitions.get(i), after.getPartition(i));
// }
}

Example 4 with BagheeraMessage

use of com.mozilla.bagheera.BagheeraProto.BagheeraMessage in project bagheera by mozilla-metrics.

the class SimpleProducer method testBasicMessage.

@Test
public void testBasicMessage() throws IOException, InterruptedException {
    // Use a ReplaySink to send messages
    String destPattern = String.format("http://localhost:%d/%s/%s/%s", BAGHEERA_PORT, SubmissionHandler.ENDPOINT_SUBMIT, TEST_NAMESPACE, ReplaySink.KEY_PLACEHOLDER);
    ReplaySink sink = new ReplaySink(destPattern, "1", "true", "true");
    sink.store(key, json.getBytes(), timestamp);
    assertEquals(1, producer.queueSize());
    BagheeraMessage message = producer.getQueue().poll();
    String payload = message.getPayload().toStringUtf8();
    assertEquals(json, payload);
    // ReplaySink doesn't preserve timestamps.
    assertNotSame(timestamp, message.getTimestamp());
    assertEquals(key, message.getId());
    assertEquals(BagheeraMessage.Operation.CREATE_UPDATE, message.getOperation());
    assertEquals("", message.getApiVersion());
    assertEquals(0, message.getPartitionCount());
}

Example 5 with BagheeraMessage

use of com.mozilla.bagheera.BagheeraProto.BagheeraMessage in project bagheera by mozilla-metrics.

the class SimpleProducer method testMessageWithPartitions.

@Test
public void testMessageWithPartitions() throws IOException, InterruptedException {
    // Use a ReplaySink to send messages
    String destPattern = String.format("http://localhost:%d/%s/%s/%s/partition1/partition2", BAGHEERA_PORT, SubmissionHandler.ENDPOINT_SUBMIT, TEST_NAMESPACE, ReplaySink.KEY_PLACEHOLDER);
    ReplaySink sink = new ReplaySink(destPattern, "1", "true", "true");
    sink.store(key, json.getBytes(), timestamp);
    assertEquals(1, producer.queueSize());
    BagheeraMessage message = producer.getQueue().poll();
    String payload = message.getPayload().toStringUtf8();
    assertEquals(json, payload);
    // ReplaySink doesn't preserve timestamps.
    assertNotSame(timestamp, message.getTimestamp());
    assertEquals(key, message.getId());
    assertEquals(BagheeraMessage.Operation.CREATE_UPDATE, message.getOperation());
    // Ensure that partition information comes through.
    assertEquals(2, message.getPartitionCount());
    assertEquals("partition1", message.getPartition(0));
    assertEquals("partition2", message.getPartition(1));
}