kafkaesque 0.0.15

Kafka implementations in Rust (Unofficial)
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121

use super::{
    lazy_broker_connection::LazyBrokerConnection, BrokerList, ClientError, MetadataClient, Result,
};
use crate::{
    clients::ClientConfig,
    models::{NodeId, PartitionId, TopicName},
};
use itertools::Itertools;
use std::{collections::HashMap, sync::Arc};
use tokio::sync::RwLock;

/// Owner of topic partition leadership information
///
/// Caches the data internally for subsequent use.
#[derive(Debug, Clone)]
pub struct TopicPartitionLeaders {
    /// Leaders of topic partitions
    cache: Arc<RwLock<HashMap<(TopicName, PartitionId), LazyBrokerConnection>>>,
    metadata_client: MetadataClient,
    config: Arc<ClientConfig>,
}

impl TopicPartitionLeaders {
    pub async fn create(config: ClientConfig) -> Result<Self> {
        let metadata_client = MetadataClient::new(config.clone());
        let cache = Default::default();
        Ok(TopicPartitionLeaders {
            cache,
            metadata_client,
            config: Arc::new(config),
        })
    }

    /// Find out and create a connection to the leader of the specified topic partition
    pub async fn get_connection_to_leader(
        &self,
        topic_name: &TopicName,
        partition_id: PartitionId,
    ) -> Result<LazyBrokerConnection> {
        let key = (topic_name.clone(), partition_id);
        let read_lock = self.cache.read().await;
        if let Some(conn) = read_lock.get(&key).cloned() {
            Ok(conn)
        } else {
            drop(read_lock);
            let topic_leaders = self.find_topic_leaders(topic_name).await?;
            let mut write_lock = self.cache.write().await;
            for (node_id, conn) in topic_leaders {
                write_lock.insert((topic_name.clone(), node_id), conn);
            }
            write_lock
                .get(&key)
                .cloned()
                .ok_or_else(|| ClientError::TopicPartitionLeaderNotFound {
                    topic_name: topic_name.clone(),
                    partition_id,
                })
        }
    }

    /// Reset the cache thus re-learning about topic partion leaders on subsequent attempts
    /// to use the connections
    pub async fn reset_cache(&self, topic_names: &[&TopicName]) -> Result<()> {
        let mut lock = self.cache.write().await;
        let keys_to_remove = lock
            .keys()
            .cloned()
            .filter(|(name, _)| topic_names.contains(&name))
            .collect_vec();
        for key in keys_to_remove {
            lock.remove(&key);
        }
        Ok(())
    }

    async fn find_topic_leaders(
        &self,
        topic_name: &TopicName,
    ) -> Result<Vec<(PartitionId, LazyBrokerConnection)>> {
        let metadata = self
            .metadata_client
            .get_metadata([topic_name.clone()])
            .await?;
        let brokers: HashMap<NodeId, LazyBrokerConnection> = metadata
            .brokers
            .into_iter()
            .map(|broker| {
                let config = ClientConfig {
                    broker_list: BrokerList::from_hostnames_and_ports([(
                        broker.host,
                        broker.port as u16,
                    )]),
                    client_id: self.config.client_id.clone(),
                };
                let conn = LazyBrokerConnection::new(config);
                (broker.id, conn)
            })
            .collect();
        let res = metadata
            .topics
            .into_iter()
            .filter(|topic| &topic.name == topic_name)
            .flat_map(|topic| topic.partitions)
            .map(|partition| {
                let conn = brokers.get(&partition.leader).cloned().ok_or(
                    ClientError::NodeMetadataNotFound {
                        node_id: partition.leader,
                    },
                )?;
                Result::Ok((partition.id, conn))
            })
            .try_collect()?;
        Ok(res)
    }
}

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
struct TopicPartition {
    pub topic_name: TopicName,
    pub partition_id: PartitionId,
}