danube_client/producer.rs
1use crate::{
2 errors::{DanubeError, Result},
3 message_router::MessageRouter,
4 retry_manager::RetryManager,
5 topic_producer::TopicProducer,
6 DanubeClient,
7};
8
9use danube_core::dispatch_strategy::ConfigDispatchStrategy;
10use danube_core::proto::schema_reference::VersionRef;
11use danube_core::proto::SchemaReference;
12use std::collections::HashMap;
13use std::sync::Arc;
14use tokio::sync::Mutex;
15use tracing::{error, info, warn};
16
17/// Represents a message producer responsible for sending messages to partitioned or non-partitioned topics distributed across message brokers.
18///
19/// The `Producer` struct is designed to handle the creation and management of a producer instance that sends messages to either partitioned or non-partitioned topics.
20/// It manages the producer's state and ensures that messages are sent according to the configured settings.
21#[derive(Debug)]
22pub struct Producer {
23 client: DanubeClient,
24 topic_name: String,
25 schema_ref: Option<SchemaReference>,
26 dispatch_strategy: ConfigDispatchStrategy,
27 producer_name: String,
28 partitions: Option<usize>,
29 message_router: Option<MessageRouter>,
30 producers: Arc<Mutex<Vec<TopicProducer>>>,
31 producer_options: ProducerOptions,
32}
33
34impl Producer {
35 pub(crate) fn new(
36 client: DanubeClient,
37 topic_name: String,
38 schema_ref: Option<SchemaReference>,
39 dispatch_strategy: Option<ConfigDispatchStrategy>,
40 producer_name: String,
41 partitions: Option<usize>,
42 message_router: Option<MessageRouter>,
43 producer_options: ProducerOptions,
44 ) -> Self {
45 let dispatch_strategy = dispatch_strategy.unwrap_or_default();
46
47 Producer {
48 client,
49 topic_name,
50 schema_ref,
51 dispatch_strategy,
52 producer_name,
53 partitions,
54 message_router,
55 producers: Arc::new(Mutex::new(Vec::new())),
56 producer_options,
57 }
58 }
59
60 /// Initializes the producer and registers it with the message brokers.
61 ///
62 /// This asynchronous method sets up the producer by establishing connections with the message brokers and configuring it for sending messages to the specified topic.
63 /// It is responsible for creating the necessary resources for producers handling partitioned topics.
64 pub async fn create(&mut self) -> Result<()> {
65 let mut topic_producers: Vec<_> = match self.partitions {
66 None => {
67 // Create a single TopicProducer for non-partitioned topic
68 vec![TopicProducer::new(
69 self.client.clone(),
70 self.topic_name.clone(),
71 self.producer_name.clone(),
72 self.schema_ref.clone(),
73 self.dispatch_strategy.clone(),
74 self.producer_options.clone(),
75 )]
76 }
77 Some(partitions) => {
78 if self.message_router.is_none() {
79 self.message_router = Some(MessageRouter::new(partitions));
80 };
81
82 (0..partitions)
83 .map(|partition_id| {
84 let topic = format!("{}-part-{}", self.topic_name, partition_id);
85 TopicProducer::new(
86 self.client.clone(),
87 topic,
88 format!("{}-{}", self.producer_name, partition_id),
89 self.schema_ref.clone(),
90 self.dispatch_strategy.clone(),
91 self.producer_options.clone(),
92 )
93 })
94 .collect()
95 }
96 };
97
98 for topic_producer in &mut topic_producers {
99 let _prod_id = topic_producer.create().await?;
100 }
101
102 // ensure that the producers are added only if all topic_producers are succesfully created
103 let mut producers = self.producers.lock().await;
104 *producers = topic_producers;
105
106 Ok(())
107 }
108
109 /// Sends a message to the topic associated with this producer.
110 ///
111 /// It handles the serialization of the payload and any user-defined attributes. This method assumes that the producer has been successfully initialized and is ready to send messages.
112 ///
113 /// # Parameters
114 ///
115 /// - `data`: The message payload to be sent. This should be a `Vec<u8>` representing the content of the message.
116 /// - `attributes`: Optional user-defined properties or attributes associated with the message. This is a `HashMap<String, String>` where keys and values represent the attribute names and values, respectively.
117 ///
118 /// # Returns
119 ///
120 /// - `Ok(u64)`: The sequence ID of the sent message if the operation is successful. This ID can be used for tracking and acknowledging the message.
121 /// - `Err(e)`: An error if message sending fails. Possible reasons for failure include network issues, serialization errors, or broker-related problems.
122 pub async fn send(
123 &self,
124 data: Vec<u8>,
125 attributes: Option<HashMap<String, String>>,
126 ) -> Result<u64> {
127 let partition = self.select_partition();
128 let retry_manager = RetryManager::new(
129 self.producer_options.max_retries,
130 self.producer_options.base_backoff_ms,
131 self.producer_options.max_backoff_ms,
132 );
133
134 let mut attempts = 0;
135
136 loop {
137 let send_result = {
138 let mut producers = self.producers.lock().await;
139 producers[partition].send(&data, attributes.as_ref(), None).await
140 };
141
142 match send_result {
143 Ok(sequence_id) => return Ok(sequence_id),
144
145 // Unrecoverable: attempt full recreation
146 Err(ref error) if matches!(error, DanubeError::Unrecoverable(_)) => {
147 warn!(error = ?error, "unrecoverable error, attempting producer recreation");
148 self.recreate_producer(partition).await?;
149 attempts = 0;
150 }
151
152 // Retryable: backoff, then escalate to lookup+recreate after max retries
153 Err(error) if retry_manager.is_retryable_error(&error) => {
154 attempts += 1;
155 if attempts > retry_manager.max_retries() {
156 warn!("max retries exceeded, attempting broker lookup and recreation");
157 self.lookup_and_recreate(partition, error).await?;
158 attempts = 0;
159 continue;
160 }
161 let backoff = retry_manager.calculate_backoff(attempts - 1);
162 tokio::time::sleep(backoff).await;
163 }
164
165 // Non-retryable: bail
166 Err(error) => {
167 error!(error = ?error, "non-retryable error in producer send");
168 return Err(error);
169 }
170 }
171 }
172 }
173
174 /// Sends a message with a routing key for Key-Shared subscriptions.
175 ///
176 /// For partitioned topics: hashes the routing key to a specific partition,
177 /// ensuring all messages with the same key go to the same partition's WAL.
178 /// For non-partitioned topics: simply tags the routing key on the message.
179 ///
180 /// All messages with the same routing key are guaranteed to be delivered
181 /// to the same consumer, in order, within a Key-Shared subscription.
182 pub async fn send_with_key(
183 &self,
184 data: Vec<u8>,
185 attributes: Option<HashMap<String, String>>,
186 routing_key: &str,
187 ) -> Result<u64> {
188 // Key-based partition selection (vs round-robin in send())
189 let partition = self.select_partition_for_key(routing_key);
190 let retry_manager = RetryManager::new(
191 self.producer_options.max_retries,
192 self.producer_options.base_backoff_ms,
193 self.producer_options.max_backoff_ms,
194 );
195
196 let mut attempts = 0;
197
198 loop {
199 let send_result = {
200 let mut producers = self.producers.lock().await;
201 producers[partition]
202 .send(&data, attributes.as_ref(), Some(routing_key))
203 .await
204 };
205
206 match send_result {
207 Ok(sequence_id) => return Ok(sequence_id),
208
209 // Unrecoverable: attempt full recreation
210 Err(ref error) if matches!(error, DanubeError::Unrecoverable(_)) => {
211 warn!(error = ?error, "unrecoverable error, attempting producer recreation");
212 self.recreate_producer(partition).await?;
213 attempts = 0;
214 }
215
216 // Retryable: backoff, then escalate to lookup+recreate after max retries
217 Err(error) if retry_manager.is_retryable_error(&error) => {
218 attempts += 1;
219 if attempts > retry_manager.max_retries() {
220 warn!("max retries exceeded, attempting broker lookup and recreation");
221 self.lookup_and_recreate(partition, error).await?;
222 attempts = 0;
223 continue;
224 }
225 let backoff = retry_manager.calculate_backoff(attempts - 1);
226 tokio::time::sleep(backoff).await;
227 }
228
229 // Non-retryable: bail
230 Err(error) => {
231 error!(error = ?error, "non-retryable error in producer send_with_key");
232 return Err(error);
233 }
234 }
235 }
236 }
237
238 /// Select the next partition using round-robin, or 0 for non-partitioned topics.
239 fn select_partition(&self) -> usize {
240 match self.partitions {
241 Some(_) => self
242 .message_router
243 .as_ref()
244 .expect("message_router must be initialized for partitioned topics")
245 .round_robin(),
246 None => 0,
247 }
248 }
249
250 /// Select partition by hashing the routing key.
251 /// Ensures all messages with the same key go to the same partition.
252 /// For non-partitioned topics, returns 0.
253 fn select_partition_for_key(&self, routing_key: &str) -> usize {
254 match self.partitions {
255 Some(_) => self
256 .message_router
257 .as_ref()
258 .expect("message_router must be initialized for partitioned topics")
259 .key_route(routing_key),
260 None => 0,
261 }
262 }
263
264 /// Recreate a single topic producer (e.g., after an unrecoverable error).
265 async fn recreate_producer(&self, partition: usize) -> Result<()> {
266 let mut producers = self.producers.lock().await;
267 producers[partition].create().await?;
268 info!("producer recreation successful");
269 Ok(())
270 }
271
272 /// Look up a new broker and recreate the topic producer on the new connection.
273 /// On lookup failure, returns the `original_error` from the failed send.
274 async fn lookup_and_recreate(
275 &self,
276 partition: usize,
277 original_error: DanubeError,
278 ) -> Result<()> {
279 let mut producers = self.producers.lock().await;
280 let producer = &mut producers[partition];
281
282 let broker_address = producer
283 .client
284 .lookup_service
285 .handle_lookup(&producer.connect_url, &producer.topic)
286 .await
287 .map_err(|_| original_error)?;
288
289 producer.broker_addr = broker_address.broker_url;
290 producer.connect_url = broker_address.connect_url;
291 producer.proxy = broker_address.proxy;
292 producer.create().await?;
293 info!("broker lookup and producer recreation successful");
294 Ok(())
295 }
296}
297
298/// A builder for creating a new `Producer` instance.
299///
300/// `ProducerBuilder` provides a fluent API for configuring and instantiating a `Producer`.
301/// It allows you to set various properties that define how the producer will behave and interact with the message broker.
302#[derive(Debug, Clone)]
303pub struct ProducerBuilder {
304 client: DanubeClient,
305 topic: Option<String>,
306 num_partitions: Option<usize>,
307 producer_name: Option<String>,
308 // TODO Phase 4: schema removed
309 // schema: Option<Schema>,
310 // Phase 5: Schema registry support
311 schema_ref: Option<SchemaReference>,
312 dispatch_strategy: Option<ConfigDispatchStrategy>,
313 producer_options: ProducerOptions,
314}
315
316impl ProducerBuilder {
317 pub fn new(client: &DanubeClient) -> Self {
318 ProducerBuilder {
319 client: client.clone(),
320 topic: None,
321 num_partitions: None,
322 producer_name: None,
323 schema_ref: None,
324 dispatch_strategy: None,
325 producer_options: ProducerOptions::default(),
326 }
327 }
328
329 /// Sets the topic name for the producer. This is a required field.
330 ///
331 /// This method specifies the topic that the producer will send messages to. It must be set before creating the producer.
332 ///
333 /// # Parameters
334 ///
335 /// - `topic`: The name of the topic for the producer. This should be a non-empty string that corresponds to an existing or new topic.
336 pub fn with_topic(mut self, topic: impl Into<String>) -> Self {
337 self.topic = Some(topic.into());
338 self
339 }
340
341 /// Sets the name of the producer. This is a required field.
342 ///
343 /// This method specifies the name to be assigned to the producer instance. It must be set before creating the producer.
344 ///
345 /// # Parameters
346 ///
347 /// - `producer_name`: The name assigned to the producer instance. This should be a non-empty string used for identifying the producer.
348 pub fn with_name(mut self, producer_name: impl Into<String>) -> Self {
349 self.producer_name = Some(producer_name.into());
350 self
351 }
352
353 // ===== Schema Registry Methods =====
354
355 /// Set schema by subject name (uses latest version)
356 ///
357 /// The producer will reference the latest schema version for the given subject.
358 /// The schema must be registered in the schema registry before use.
359 ///
360 /// # Example
361 /// ```no_run
362 /// # use danube_client::DanubeClient;
363 /// # async fn example(client: DanubeClient) -> Result<(), Box<dyn std::error::Error>> {
364 /// let mut producer = client.new_producer()
365 /// .with_topic("user-events")
366 /// .with_name("my-producer")
367 /// .with_schema_subject("user-events-value") // Uses latest version
368 /// .build()?;
369 /// # Ok(())
370 /// # }
371 /// ```
372 pub fn with_schema_subject(mut self, subject: impl Into<String>) -> Self {
373 self.schema_ref = Some(SchemaReference {
374 subject: subject.into(),
375 version_ref: Some(VersionRef::UseLatest(true)),
376 });
377 self
378 }
379
380 /// Set schema with a pinned version
381 ///
382 /// The producer will use a specific schema version and won't automatically
383 /// upgrade to newer versions.
384 ///
385 /// # Example
386 /// ```no_run
387 /// # use danube_client::DanubeClient;
388 /// # async fn example(client: DanubeClient) -> Result<(), Box<dyn std::error::Error>> {
389 /// let mut producer = client.new_producer()
390 /// .with_topic("user-events")
391 /// .with_name("my-producer")
392 /// .with_schema_version("user-events-value", 2) // Pin to version 2
393 /// .build()?;
394 /// # Ok(())
395 /// # }
396 /// ```
397 pub fn with_schema_version(mut self, subject: impl Into<String>, version: u32) -> Self {
398 self.schema_ref = Some(SchemaReference {
399 subject: subject.into(),
400 version_ref: Some(VersionRef::PinnedVersion(version)),
401 });
402 self
403 }
404
405 /// Set schema with a minimum version requirement
406 ///
407 /// The producer will use the specified version or any newer compatible version.
408 ///
409 /// # Example
410 /// ```no_run
411 /// # use danube_client::DanubeClient;
412 /// # async fn example(client: DanubeClient) -> Result<(), Box<dyn std::error::Error>> {
413 /// let mut producer = client.new_producer()
414 /// .with_topic("user-events")
415 /// .with_name("my-producer")
416 /// .with_schema_min_version("user-events-value", 2) // Use v2 or newer
417 /// .build()?;
418 /// # Ok(())
419 /// # }
420 /// ```
421 pub fn with_schema_min_version(mut self, subject: impl Into<String>, min_version: u32) -> Self {
422 self.schema_ref = Some(SchemaReference {
423 subject: subject.into(),
424 version_ref: Some(VersionRef::MinVersion(min_version)),
425 });
426 self
427 }
428
429 /// Set schema with a custom SchemaReference (advanced use)
430 ///
431 /// This allows full control over schema versioning. For most use cases,
432 /// prefer `with_schema_subject()`, `with_schema_version()`, or `with_schema_min_version()`.
433 ///
434 /// # Example
435 /// ```no_run
436 /// # use danube_client::{DanubeClient, SchemaReference, VersionRef};
437 /// # async fn example(client: DanubeClient) -> Result<(), Box<dyn std::error::Error>> {
438 /// let mut producer = client.new_producer()
439 /// .with_topic("user-events")
440 /// .with_name("my-producer")
441 /// .with_schema_reference(SchemaReference {
442 /// subject: "user-events-value".to_string(),
443 /// version_ref: Some(VersionRef::PinnedVersion(2)),
444 /// })
445 /// .build()?;
446 /// # Ok(())
447 /// # }
448 /// ```
449 pub fn with_schema_reference(mut self, schema_ref: SchemaReference) -> Self {
450 self.schema_ref = Some(schema_ref);
451 self
452 }
453
454 /// Sets the reliable dispatch options for the producer.
455 /// This method configures the dispatch strategy for the producer, which determines how messages are stored and managed.
456 /// The dispatch strategy defines how long messages are retained and how they are managed in the message broker.
457 ///
458 /// # Parameters
459 ///
460 /// No parameters; broker uses defaults for reliable topics.
461 pub fn with_reliable_dispatch(mut self) -> Self {
462 let dispatch_strategy = ConfigDispatchStrategy::Reliable;
463 self.dispatch_strategy = Some(dispatch_strategy);
464 self
465 }
466
467 /// Sets the configuration options for the producer, allowing customization of producer behavior.
468 ///
469 /// This method allows you to specify various configuration options that affect how the producer operates.
470 /// These options can control aspects such as retries, timeouts, and other producer-specific settings.
471 ///
472 /// # Parameters
473 ///
474 /// - `options`: A `ProducerOptions` instance containing the configuration options for the producer. This should be configured according to the desired behavior and requirements of the producer.
475 pub fn with_options(mut self, options: ProducerOptions) -> Self {
476 self.producer_options = options;
477 self
478 }
479
480 /// Sets the number of partitions for the topic.
481 ///
482 /// This method specifies how many partitions the topic should have. Partitions are used to distribute the load of messages across multiple Danube brokers, which can help with parallel processing and scalability.
483 ///
484 /// # Parameters
485 ///
486 /// - `partitions`: The number of partitions for the topic. This should be a positive integer representing the desired number of partitions. More partitions can improve parallelism and throughput. Default is 0 = non-partitioned topic.
487 pub fn with_partitions(mut self, partitions: usize) -> Self {
488 self.num_partitions = Some(partitions);
489 self
490 }
491
492 /// Creates a new `Producer` instance using the settings configured in the `ProducerBuilder`.
493 ///
494 /// This method performs validation to ensure that all required fields are set before creating the `Producer`. Once validation is successful, it constructs and returns a new `Producer` instance configured with the specified settings.
495 ///
496 /// # Returns
497 ///
498 /// - A `Producer` instance if the builder configuration is valid and the producer is created successfully.
499 ///
500 /// # Example
501 /// ```no_run
502 /// # use danube_client::DanubeClient;
503 /// # async fn example(client: DanubeClient) -> Result<(), Box<dyn std::error::Error>> {
504 /// let mut producer = client.new_producer()
505 /// .with_topic("my-topic")
506 /// .with_name("my-producer")
507 /// .with_partitions(3)
508 /// .build()?;
509 /// # Ok(())
510 /// # }
511 /// ```
512 pub fn build(self) -> Result<Producer> {
513 let topic_name = self.topic.ok_or_else(|| {
514 DanubeError::Unrecoverable("topic is required to build a Producer".into())
515 })?;
516 let producer_name = self.producer_name.ok_or_else(|| {
517 DanubeError::Unrecoverable("producer name is required to build a Producer".into())
518 })?;
519
520 if let Some(0) = self.num_partitions {
521 return Err(DanubeError::Unrecoverable(
522 "partitions must be > 0 or omitted for non-partitioned topic".into(),
523 ));
524 }
525
526 Ok(Producer::new(
527 self.client,
528 topic_name,
529 self.schema_ref,
530 self.dispatch_strategy,
531 producer_name,
532 self.num_partitions,
533 None,
534 self.producer_options,
535 ))
536 }
537}
538
539/// Configuration options for producers
540#[derive(Debug, Clone, Default)]
541#[non_exhaustive]
542pub struct ProducerOptions {
543 // Maximum number of retries for operations like create/send on transient failures
544 pub max_retries: usize,
545 // Base backoff in milliseconds for exponential backoff
546 pub base_backoff_ms: u64,
547 // Maximum backoff cap in milliseconds
548 pub max_backoff_ms: u64,
549}
550
551impl ProducerOptions {
552 /// Create new ProducerOptions with explicit retry settings.
553 pub fn new(max_retries: usize, base_backoff_ms: u64, max_backoff_ms: u64) -> Self {
554 Self {
555 max_retries,
556 base_backoff_ms,
557 max_backoff_ms,
558 }
559 }
560}