protected Optional<String> getRoutingKey(OutgoingMessageEnvelope envelope) { Object partitionKey = envelope.getPartitionKey(); if (partitionKey == null) { return Optional.absent(); } return Optional.of(partitionKey.toString()); }
protected Object getEnvelopePartitionId(OutgoingMessageEnvelope envelope) { return envelope.getPartitionKey() == null ? envelope.getKey() : envelope.getPartitionKey(); }
/** * Sends a specified message envelope from a specified Samza source. * @param source String representing the source of the message. * @param envelope Aggregate object representing the serialized message to send from the source. */ @Override public void send(String source, OutgoingMessageEnvelope envelope) { Object key = envelope.getKey(); Object message = envelope.getMessage(); Object partitionKey; // We use the partition key from message if available, if not fallback to message key or use message as partition // key as the final resort. if (envelope.getPartitionKey() != null) { partitionKey = envelope.getPartitionKey(); } else if (key != null) { partitionKey = key; } else { partitionKey = message; } Preconditions.checkNotNull(partitionKey, "Failed to compute partition key for the message: " + envelope); int partition = Math.abs(partitionKey.hashCode()) % memoryManager.getPartitionCountForSystemStream(envelope.getSystemStream()); SystemStreamPartition ssp = new SystemStreamPartition(envelope.getSystemStream(), new Partition(partition)); memoryManager.put(ssp, key, message); }
/** * Sends a specified message envelope from a specified Samza source. * @param source String representing the source of the message. * @param envelope Aggregate object representing the serialized message to send from the source. */ @Override public void send(String source, OutgoingMessageEnvelope envelope) { Object key = envelope.getKey(); Object message = envelope.getMessage(); Object partitionKey; // We use the partition key from message if available, if not fallback to message key or use message as partition // key as the final resort. if (envelope.getPartitionKey() != null) { partitionKey = envelope.getPartitionKey(); } else if (key != null) { partitionKey = key; } else { partitionKey = message; } Preconditions.checkNotNull(partitionKey, "Failed to compute partition key for the message: " + envelope); int partition = Math.abs(partitionKey.hashCode()) % memoryManager.getPartitionCountForSystemStream(envelope.getSystemStream()); SystemStreamPartition ssp = new SystemStreamPartition(envelope.getSystemStream(), new Partition(partition)); memoryManager.put(ssp, key, message); }
/** * Sends a specified message envelope from a specified Samza source. * @param source String representing the source of the message. * @param envelope Aggregate object representing the serialized message to send from the source. */ @Override public void send(String source, OutgoingMessageEnvelope envelope) { Object key = envelope.getKey(); Object message = envelope.getMessage(); Object partitionKey; // We use the partition key from message if available, if not fallback to message key or use message as partition // key as the final resort. if (envelope.getPartitionKey() != null) { partitionKey = envelope.getPartitionKey(); } else if (key != null) { partitionKey = key; } else { partitionKey = message; } Preconditions.checkNotNull(partitionKey, "Failed to compute partition key for the message: " + envelope); int partition = Math.abs(partitionKey.hashCode()) % memoryManager.getPartitionCountForSystemStream(envelope.getSystemStream()); SystemStreamPartition ssp = new SystemStreamPartition(envelope.getSystemStream(), new Partition(partition)); memoryManager.put(ssp, key, message); }
/** * Sends a specified message envelope from a specified Samza source. * @param source String representing the source of the message. * @param envelope Aggregate object representing the serialized message to send from the source. */ @Override public void send(String source, OutgoingMessageEnvelope envelope) { Object key = envelope.getKey(); Object message = envelope.getMessage(); Object partitionKey; // We use the partition key from message if available, if not fallback to message key or use message as partition // key as the final resort. if (envelope.getPartitionKey() != null) { partitionKey = envelope.getPartitionKey(); } else if (key != null) { partitionKey = key; } else { partitionKey = message; } Preconditions.checkNotNull(partitionKey, "Failed to compute partition key for the message: " + envelope); int partition = Math.abs(partitionKey.hashCode()) % memoryManager.getPartitionCountForSystemStream(envelope.getSystemStream()); SystemStreamPartition ssp = new SystemStreamPartition(envelope.getSystemStream(), new Partition(partition)); memoryManager.put(ssp, key, message); }
/** * Sends a specified message envelope from a specified Samza source. * @param source String representing the source of the message. * @param envelope Aggregate object representing the serialized message to send from the source. */ @Override public void send(String source, OutgoingMessageEnvelope envelope) { Object key = envelope.getKey(); Object message = envelope.getMessage(); Object partitionKey; // We use the partition key from message if available, if not fallback to message key or use message as partition // key as the final resort. if (envelope.getPartitionKey() != null) { partitionKey = envelope.getPartitionKey(); } else if (key != null) { partitionKey = key; } else { partitionKey = message; } Preconditions.checkNotNull(partitionKey, "Failed to compute partition key for the message: " + envelope); int partition = Math.abs(partitionKey.hashCode()) % memoryManager.getPartitionCountForSystemStream(envelope.getSystemStream()); SystemStreamPartition ssp = new SystemStreamPartition(envelope.getSystemStream(), new Partition(partition)); memoryManager.put(ssp, key, message); }
@Override public void send(String source, OutgoingMessageEnvelope envelope) { SystemStream systemStream = envelope.getSystemStream(); List<SystemStreamPartition> sspForSystem = MSG_QUEUES.keySet().stream() .filter(ssp -> ssp.getSystemStream().equals(systemStream)) .collect(ArrayList::new, (l, ssp) -> l.add(ssp), (l1, l2) -> l1.addAll(l2)); if (sspForSystem.isEmpty()) { MSG_QUEUES.putIfAbsent(new SystemStreamPartition(systemStream, new Partition(0)), new ArrayList<>()); sspForSystem.add(new SystemStreamPartition(systemStream, new Partition(0))); } int partitionCount = sspForSystem.size(); int partitionId = envelope.getPartitionKey() == null ? envelope.getKey() == null ? this.seed.nextInt(partitionCount) : envelope.getKey().hashCode() % partitionCount : envelope.getPartitionKey().hashCode() % partitionCount; SystemStreamPartition ssp = new SystemStreamPartition(envelope.getSystemStream(), new Partition(partitionId)); List<IncomingMessageEnvelope> msgQueue = MSG_QUEUES.get(ssp); msgQueue.add(new IncomingMessageEnvelope(ssp, null, envelope.getKey(), envelope.getMessage())); }
@Test public void testSend() { SystemStreamMetadata metadata = mock(SystemStreamMetadata.class); Map<Partition, SystemStreamMetadata.SystemStreamPartitionMetadata> partitionMetadata = new HashMap<>(); partitionMetadata.put(new Partition(0), mock(SystemStreamMetadata.SystemStreamPartitionMetadata.class)); partitionMetadata.put(new Partition(1), mock(SystemStreamMetadata.SystemStreamPartitionMetadata.class)); partitionMetadata.put(new Partition(2), mock(SystemStreamMetadata.SystemStreamPartitionMetadata.class)); partitionMetadata.put(new Partition(3), mock(SystemStreamMetadata.SystemStreamPartitionMetadata.class)); when(metadata.getSystemStreamPartitionMetadata()).thenReturn(partitionMetadata); StreamMetadataCache metadataCache = mock(StreamMetadataCache.class); when(metadataCache.getSystemStreamMetadata(anyObject(), anyBoolean())).thenReturn(metadata); SystemStream systemStream = new SystemStream("test-system", "test-stream"); Set<Integer> partitions = new HashSet<>(); MessageCollector collector = mock(MessageCollector.class); doAnswer(invocation -> { OutgoingMessageEnvelope envelope = (OutgoingMessageEnvelope) invocation.getArguments()[0]; partitions.add((Integer) envelope.getPartitionKey()); assertEquals(envelope.getSystemStream(), systemStream); return null; }).when(collector).send(any()); ControlMessageSender sender = new ControlMessageSender(metadataCache); WatermarkMessage watermark = new WatermarkMessage(System.currentTimeMillis(), "task 0"); sender.send(watermark, systemStream, collector); assertEquals(partitions.size(), 1); }
@Test public void testBroadcast() { SystemStreamMetadata metadata = mock(SystemStreamMetadata.class); Map<Partition, SystemStreamMetadata.SystemStreamPartitionMetadata> partitionMetadata = new HashMap<>(); partitionMetadata.put(new Partition(0), mock(SystemStreamMetadata.SystemStreamPartitionMetadata.class)); partitionMetadata.put(new Partition(1), mock(SystemStreamMetadata.SystemStreamPartitionMetadata.class)); partitionMetadata.put(new Partition(2), mock(SystemStreamMetadata.SystemStreamPartitionMetadata.class)); partitionMetadata.put(new Partition(3), mock(SystemStreamMetadata.SystemStreamPartitionMetadata.class)); when(metadata.getSystemStreamPartitionMetadata()).thenReturn(partitionMetadata); StreamMetadataCache metadataCache = mock(StreamMetadataCache.class); when(metadataCache.getSystemStreamMetadata(anyObject(), anyBoolean())).thenReturn(metadata); SystemStream systemStream = new SystemStream("test-system", "test-stream"); Set<Integer> partitions = new HashSet<>(); MessageCollector collector = mock(MessageCollector.class); doAnswer(invocation -> { OutgoingMessageEnvelope envelope = (OutgoingMessageEnvelope) invocation.getArguments()[0]; partitions.add((Integer) envelope.getPartitionKey()); assertEquals(envelope.getSystemStream(), systemStream); return null; }).when(collector).send(any()); ControlMessageSender sender = new ControlMessageSender(metadataCache); WatermarkMessage watermark = new WatermarkMessage(System.currentTimeMillis(), "task 0"); SystemStreamPartition ssp = new SystemStreamPartition(systemStream, new Partition(0)); sender.broadcastToOtherPartitions(watermark, ssp, collector); assertEquals(partitions.size(), 3); } }