kf-socket 0.1.0

Provide TCP socket wrapper using kafka protocol
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
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142


    

use std::io::Error as IoError;
use std::env::temp_dir;
use std::net::SocketAddr;
use std::time::Duration;

use log::debug;
use futures::io::AsyncWriteExt;
use futures::future::join;
use futures::stream::StreamExt;

use flv_future_core::test_async;
use flv_future_core::sleep;
use flv_future_aio::fs::file_util;
use flv_future_aio::fs::AsyncFile;
use flv_future_aio::net::AsyncTcpListener;
use kf_protocol::Encoder;
use kf_protocol::api::Request;
use kf_protocol::api::ResponseMessage;
use kf_protocol::api::RequestMessage;
use kf_protocol::api::DefaultBatch;
use kf_protocol::api::DefaultRecord;
use kf_protocol::message::fetch::DefaultKfFetchRequest;
use kf_socket::KfSocketError;
use kf_protocol::fs::FileFetchResponse;
use kf_protocol::fs::KfFileFetchRequest;
use kf_protocol::fs::FilePartitionResponse;
use kf_protocol::fs::FileTopicResponse;

use flv_util::fixture::ensure_clean_file;
use kf_socket::KfSocket;


/// create sample batches with message
fn create_batches(records: u16) -> DefaultBatch {
    let mut batches = DefaultBatch::default();
    let header = batches.get_mut_header();
    header.magic = 2;
    header.producer_id = 20;
    header.producer_epoch = -1;

    for i in 0..records {
        let msg = format!("record {}", i);
        let record: DefaultRecord = msg.into();
        batches.add_record(record);
    }
    batches
}

async fn setup_batch_file() -> Result<(), IoError> {
    let test_file_path = temp_dir().join("batch_fetch");
    ensure_clean_file(&test_file_path);
    debug!("creating test file: {:#?}", test_file_path);
    let mut file = file_util::create(&test_file_path).await?;
    let batch = create_batches(2);
    let bytes = batch.as_bytes(0)?;
    file.write_all(bytes.as_ref()).await?;
    Ok(())
}

async fn test_server(addr: SocketAddr) -> Result<(), KfSocketError> {
    let listener = AsyncTcpListener::bind(&addr).await?;
    debug!("server is running");
    let mut incoming = listener.incoming();
    let incoming_stream = incoming.next().await;
    debug!("server: got connection");
    let incoming_stream = incoming_stream.expect("next").expect("unwrap again");
    let mut socket: KfSocket = incoming_stream.into();

    let fetch_request: Result<RequestMessage<KfFileFetchRequest>, KfSocketError> = socket
        .get_mut_stream()
        .next_request_item()
        .await
        .expect("next value");
    let request = fetch_request?;
    debug!("received fetch request: {:#?}", request);

    let test_file_path = temp_dir().join("batch_fetch");
    debug!("opening file test file: {:#?}", test_file_path);
    let file = file_util::open(&test_file_path).await?;

    let mut response = FileFetchResponse::default();
    let mut topic_response = FileTopicResponse::default();
    let mut part_response = FilePartitionResponse::default();
    part_response.partition_index = 10;
    part_response.records = file.as_slice(0, None).await?.into();
    topic_response.partitions.push(part_response);
    response.topics.push(topic_response);
    let resp_msg = ResponseMessage::new(10, response);

    debug!(
        "response message write size: {}",
        resp_msg.write_size(KfFileFetchRequest::DEFAULT_API_VERSION)
    );

    socket
        .get_mut_sink()
        .encode_file_slices(&resp_msg, KfFileFetchRequest::DEFAULT_API_VERSION)
        .await?;
    debug!("server: finish sending out");
    Ok(())
}

async fn setup_client(addr: SocketAddr) -> Result<(), KfSocketError> {
    sleep(Duration::from_millis(50)).await;
    debug!("client: trying to connect");
    let mut socket = KfSocket::connect(&addr).await?;
    debug!("client: connect to test server and waiting...");

    let req_msg: RequestMessage<DefaultKfFetchRequest> = RequestMessage::default();
    let res_msg = socket.send(&req_msg).await?;

    debug!("output: {:#?}", res_msg);
    let topic_responses = res_msg.response.topics;
    assert_eq!(topic_responses.len(), 1);
    let part_responses = &topic_responses[0].partitions;
    assert_eq!(part_responses.len(), 1);
    let batches = &part_responses[0].records.batches;
    assert_eq!(batches.len(), 1);
    let records = &batches[0].records;
    assert_eq!(records.len(), 2);
    assert_eq!(records[0].value.to_string(), "record 0");
    assert_eq!(records[1].value.to_string(), "record 1");

    Ok(())
}

/// test server where it is sending out file copy
#[test_async]
async fn test_save_fetch() -> Result<(), KfSocketError> {
    // create fetch and save
    setup_batch_file().await?;

    let addr = "127.0.0.1:9911".parse::<SocketAddr>().expect("parse");

    let _r = join(setup_client(addr), test_server(addr.clone())).await;

    Ok(())
}