Skip to content

Commit

Permalink
[ISSUE #1012]🚀Support client Broadcasting consume-local file store⚡️
Browse files Browse the repository at this point in the history
  • Loading branch information
mxsm committed Sep 28, 2024
1 parent 5cefbc9 commit 8c74161
Show file tree
Hide file tree
Showing 4 changed files with 230 additions and 12 deletions.
1 change: 1 addition & 0 deletions rocketmq-client/src/consumer/store.rs
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@

mod controllable_offset;
pub(crate) mod local_file_offset_store;
mod offset_serialize_wrapper;
pub(crate) mod offset_store;
pub(crate) mod read_offset_type;
pub(crate) mod remote_broker_offset_store;
7 changes: 7 additions & 0 deletions rocketmq-client/src/consumer/store/controllable_offset.rs
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,13 @@ impl ControllableOffset {
}
}

pub fn new_atomic(value: AtomicI64) -> Self {
Self {
value: Arc::new(value),
allow_to_update: Arc::new(AtomicBool::new(true)),
}
}

Check warning on line 40 in rocketmq-client/src/consumer/store/controllable_offset.rs

View check run for this annotation

Codecov / codecov/patch

rocketmq-client/src/consumer/store/controllable_offset.rs#L35-L40

Added lines #L35 - L40 were not covered by tests

pub fn update(&self, target: i64, increase_only: bool) {
if self.allow_to_update.load(Ordering::SeqCst) {
self.value
Expand Down
206 changes: 194 additions & 12 deletions rocketmq-client/src/consumer/store/local_file_offset_store.rs
Original file line number Diff line number Diff line change
Expand Up @@ -16,53 +16,235 @@
*/
use std::collections::HashMap;
use std::collections::HashSet;
use std::path::PathBuf;
use std::sync::atomic::AtomicI64;
use std::sync::atomic::Ordering;
use std::sync::Arc;

use once_cell::sync::Lazy;
use rocketmq_common::common::message::message_queue::MessageQueue;
use rocketmq_common::utils::file_utils;
use rocketmq_common::ArcRefCellWrapper;
use rocketmq_remoting::protocol::RemotingDeserializable;
use rocketmq_remoting::protocol::RemotingSerializable;
use tokio::sync::Mutex;
use tracing::error;
use tracing::info;

use crate::consumer::store::controllable_offset::ControllableOffset;
use crate::consumer::store::offset_serialize_wrapper::OffsetSerializeWrapper;
use crate::consumer::store::offset_store::OffsetStoreTrait;
use crate::consumer::store::read_offset_type::ReadOffsetType;
use crate::error::MQClientError;
use crate::factory::mq_client_instance::MQClientInstance;
use crate::Result;

pub struct LocalFileOffsetStore;
static LOCAL_OFFSET_STORE_DIR: Lazy<PathBuf> = Lazy::new(|| {

Check warning on line 42 in rocketmq-client/src/consumer/store/local_file_offset_store.rs

View check run for this annotation

Codecov / codecov/patch

rocketmq-client/src/consumer/store/local_file_offset_store.rs#L42

Added line #L42 was not covered by tests
#[cfg(target_os = "windows")]
let home = std::env::var("user.home")
.map_or(PathBuf::from("C:\\tmp\\.rocketmq_offsets"), |home| {
PathBuf::from(home).join(".rocketmq_offsets")
});

#[cfg(not(target_os = "windows"))]
let home = std::env::var("user.home").map_or(PathBuf::from("/tmp/.rocketmq_offsets"), |home| {
PathBuf::from(home).join(".rocketmq_offsets")
});

Check warning on line 52 in rocketmq-client/src/consumer/store/local_file_offset_store.rs

View check run for this annotation

Codecov / codecov/patch

rocketmq-client/src/consumer/store/local_file_offset_store.rs#L50-L52

Added lines #L50 - L52 were not covered by tests

std::env::var("rocketmq.client.localOffsetStoreDir").map_or(home, PathBuf::from)
});

Check warning on line 55 in rocketmq-client/src/consumer/store/local_file_offset_store.rs

View check run for this annotation

Codecov / codecov/patch

rocketmq-client/src/consumer/store/local_file_offset_store.rs#L54-L55

Added lines #L54 - L55 were not covered by tests

pub struct LocalFileOffsetStore {
client_instance: ArcRefCellWrapper<MQClientInstance>,
group_name: String,
store_path: String,
offset_table: Arc<Mutex<HashMap<MessageQueue, ControllableOffset>>>,
}

impl LocalFileOffsetStore {
pub fn new(mq_client_factory: ArcRefCellWrapper<MQClientInstance>, group_name: String) -> Self {
Self
pub fn new(client_instance: ArcRefCellWrapper<MQClientInstance>, group_name: String) -> Self {
Self {
client_instance,
group_name,
store_path: LOCAL_OFFSET_STORE_DIR

Check warning on line 69 in rocketmq-client/src/consumer/store/local_file_offset_store.rs

View check run for this annotation

Codecov / codecov/patch

rocketmq-client/src/consumer/store/local_file_offset_store.rs#L65-L69

Added lines #L65 - L69 were not covered by tests
.clone()
.join("offsets.json")
.to_string_lossy()
.to_string(),
offset_table: Arc::new(Mutex::new(HashMap::new())),
}
}

Check warning on line 76 in rocketmq-client/src/consumer/store/local_file_offset_store.rs

View check run for this annotation

Codecov / codecov/patch

rocketmq-client/src/consumer/store/local_file_offset_store.rs#L74-L76

Added lines #L74 - L76 were not covered by tests

fn read_local_offset(&self) -> Result<Option<OffsetSerializeWrapper>> {
let content =
file_utils::file_to_string(&self.store_path).map_or("".to_string(), |content| content);
if content.is_empty() {
self.read_local_offset_bak()

Check warning on line 82 in rocketmq-client/src/consumer/store/local_file_offset_store.rs

View check run for this annotation

Codecov / codecov/patch

rocketmq-client/src/consumer/store/local_file_offset_store.rs#L78-L82

Added lines #L78 - L82 were not covered by tests
} else {
match OffsetSerializeWrapper::decode(content.as_bytes()) {
Ok(value) => Ok(Some(value)),
Err(_) => Err(MQClientError::MQClientErr(

Check warning on line 86 in rocketmq-client/src/consumer/store/local_file_offset_store.rs

View check run for this annotation

Codecov / codecov/patch

rocketmq-client/src/consumer/store/local_file_offset_store.rs#L84-L86

Added lines #L84 - L86 were not covered by tests
-1,
format!("read local offset failed, content: {}", content),
)),

Check warning on line 89 in rocketmq-client/src/consumer/store/local_file_offset_store.rs

View check run for this annotation

Codecov / codecov/patch

rocketmq-client/src/consumer/store/local_file_offset_store.rs#L88-L89

Added lines #L88 - L89 were not covered by tests
}
}
}
fn read_local_offset_bak(&self) -> Result<Option<OffsetSerializeWrapper>> {
let content = file_utils::file_to_string(&format!("{}{}", self.store_path, ".bak"))
.map_or("".to_string(), |content| content);
if content.is_empty() {
Ok(None)

Check warning on line 97 in rocketmq-client/src/consumer/store/local_file_offset_store.rs

View check run for this annotation

Codecov / codecov/patch

rocketmq-client/src/consumer/store/local_file_offset_store.rs#L91-L97

Added lines #L91 - L97 were not covered by tests
} else {
match OffsetSerializeWrapper::decode(content.as_bytes()) {
Ok(value) => Ok(Some(value)),
Err(_) => Err(MQClientError::MQClientErr(

Check warning on line 101 in rocketmq-client/src/consumer/store/local_file_offset_store.rs

View check run for this annotation

Codecov / codecov/patch

rocketmq-client/src/consumer/store/local_file_offset_store.rs#L99-L101

Added lines #L99 - L101 were not covered by tests
-1,
format!("read local offset bak failed, content: {}", content),
)),

Check warning on line 104 in rocketmq-client/src/consumer/store/local_file_offset_store.rs

View check run for this annotation

Codecov / codecov/patch

rocketmq-client/src/consumer/store/local_file_offset_store.rs#L103-L104

Added lines #L103 - L104 were not covered by tests
}
}

Check warning on line 106 in rocketmq-client/src/consumer/store/local_file_offset_store.rs

View check run for this annotation

Codecov / codecov/patch

rocketmq-client/src/consumer/store/local_file_offset_store.rs#L106

Added line #L106 was not covered by tests
}
}

impl OffsetStoreTrait for LocalFileOffsetStore {
async fn load(&self) -> crate::Result<()> {
todo!()
let offset_serialize_wrapper = self.read_local_offset()?;
if let Some(offset_serialize_wrapper) = offset_serialize_wrapper {
let offset_table = offset_serialize_wrapper.offset_table;
let mut offset_table_inner = self.offset_table.lock().await;
for (mq, offset) in offset_table {
let offset = offset.load(Ordering::Relaxed);
info!(
"load consumer's offset, {} {} {}",
self.group_name, mq, offset
);
offset_table_inner.insert(mq, ControllableOffset::new(offset));
}
}
Ok(())
}

async fn update_offset(&self, mq: &MessageQueue, offset: i64, increase_only: bool) {
todo!()
let mut offset_table = self.offset_table.lock().await;
let offset_old = offset_table
.entry(mq.clone())
.or_insert_with(|| ControllableOffset::new(offset));
if increase_only {
offset_old.update(offset, true);

Check warning on line 134 in rocketmq-client/src/consumer/store/local_file_offset_store.rs

View check run for this annotation

Codecov / codecov/patch

rocketmq-client/src/consumer/store/local_file_offset_store.rs#L129-L134

Added lines #L129 - L134 were not covered by tests
} else {
offset_old.update_unconditionally(offset);

Check warning on line 136 in rocketmq-client/src/consumer/store/local_file_offset_store.rs

View check run for this annotation

Codecov / codecov/patch

rocketmq-client/src/consumer/store/local_file_offset_store.rs#L136

Added line #L136 was not covered by tests
}
}

async fn update_and_freeze_offset(&self, mq: &MessageQueue, offset: i64) {
todo!()
let mut offset_table = self.offset_table.lock().await;
offset_table
.entry(mq.clone())
.or_insert_with(|| ControllableOffset::new(offset))
.update_and_freeze(offset);
}

async fn read_offset(&self, mq: &MessageQueue, type_: ReadOffsetType) -> i64 {
todo!()
match type_ {
ReadOffsetType::ReadFromMemory | ReadOffsetType::MemoryFirstThenStore => {
let offset_table = self.offset_table.lock().await;
if let Some(offset) = offset_table.get(mq) {
offset.get_offset()
} else {
-1
}
}
ReadOffsetType::ReadFromStore => match self.read_local_offset() {
Ok(offset_serialize_wrapper) => {
if let Some(offset_serialize_wrapper) = offset_serialize_wrapper {
if let Some(offset) = offset_serialize_wrapper.offset_table.get(mq) {
offset.load(Ordering::Relaxed)
} else {
-1
}
} else {
-1
}
}
Err(_) => -1,
},
}
}

async fn persist_all(&mut self, mqs: &HashSet<MessageQueue>) {
todo!()
if mqs.is_empty() {

Check warning on line 176 in rocketmq-client/src/consumer/store/local_file_offset_store.rs

View check run for this annotation

Codecov / codecov/patch

rocketmq-client/src/consumer/store/local_file_offset_store.rs#L176

Added line #L176 was not covered by tests
return;
}
let mut offset_serialize_wrapper = match self.read_local_offset() {
Ok(value) => value.unwrap_or_default(),
Err(e) => {
error!("read local offset failed: {}", e);

Check warning on line 182 in rocketmq-client/src/consumer/store/local_file_offset_store.rs

View check run for this annotation

Codecov / codecov/patch

rocketmq-client/src/consumer/store/local_file_offset_store.rs#L179-L182

Added lines #L179 - L182 were not covered by tests
return;
}

Check warning on line 184 in rocketmq-client/src/consumer/store/local_file_offset_store.rs

View check run for this annotation

Codecov / codecov/patch

rocketmq-client/src/consumer/store/local_file_offset_store.rs#L184

Added line #L184 was not covered by tests
};
let offset_table = self.offset_table.lock().await;
for (mq, offset) in offset_table.iter() {
if mqs.contains(mq) {
offset_serialize_wrapper

Check warning on line 189 in rocketmq-client/src/consumer/store/local_file_offset_store.rs

View check run for this annotation

Codecov / codecov/patch

rocketmq-client/src/consumer/store/local_file_offset_store.rs#L186-L189

Added lines #L186 - L189 were not covered by tests
.offset_table
.insert(mq.clone(), AtomicI64::new(offset.get_offset()));

Check warning on line 191 in rocketmq-client/src/consumer/store/local_file_offset_store.rs

View check run for this annotation

Codecov / codecov/patch

rocketmq-client/src/consumer/store/local_file_offset_store.rs#L191

Added line #L191 was not covered by tests
}
}
let content = offset_serialize_wrapper.to_json_pretty();
if !content.is_empty() {
if let Err(e) = file_utils::string_to_file(&content, &self.store_path) {
error!(

Check warning on line 197 in rocketmq-client/src/consumer/store/local_file_offset_store.rs

View check run for this annotation

Codecov / codecov/patch

rocketmq-client/src/consumer/store/local_file_offset_store.rs#L194-L197

Added lines #L194 - L197 were not covered by tests
"persistAll consumer offset Exception, {},{}",
self.store_path, e
);
}
}

Check warning on line 202 in rocketmq-client/src/consumer/store/local_file_offset_store.rs

View check run for this annotation

Codecov / codecov/patch

rocketmq-client/src/consumer/store/local_file_offset_store.rs#L201-L202

Added lines #L201 - L202 were not covered by tests
}

async fn persist(&mut self, mq: &MessageQueue) {
todo!()
let offset_table = self.offset_table.lock().await;
if let Some(offset) = offset_table.get(mq) {
let mut offset_serialize_wrapper = match self.read_local_offset() {
Ok(value) => value.unwrap_or_default(),
Err(e) => {
error!("read local offset failed: {}", e);
return;
}
};
offset_serialize_wrapper
.offset_table
.insert(mq.clone(), AtomicI64::new(offset.get_offset()));
let content = offset_serialize_wrapper.to_json_pretty();
if !content.is_empty() {
if let Err(e) = file_utils::string_to_file(&content, &self.store_path) {
error!(
"persist consumer offset Exception, {},{}",
self.store_path, e
);
}
}
}
}

async fn remove_offset(&self, mq: &MessageQueue) {
todo!()
let mut offset_table = self.offset_table.lock().await;
offset_table.remove(mq);
info!(
"remove unnecessary messageQueue offset. group={}, mq={}, offsetTableSize={}",
mq,
self.group_name,
offset_table.len()
);
}

async fn clone_offset_table(&self, topic: &str) -> HashMap<MessageQueue, i64> {
todo!()
let offset_table = self.offset_table.lock().await;
offset_table
.iter()
.filter(|(mq, _)| topic.is_empty() || mq.get_topic() == topic)
.map(|(mq, offset)| (mq.clone(), offset.get_offset()))
.collect()
}

async fn update_consume_offset_to_broker(
Expand All @@ -71,6 +253,6 @@ impl OffsetStoreTrait for LocalFileOffsetStore {
offset: i64,
is_oneway: bool,
) -> crate::Result<()> {
todo!()
Ok(())
}
}
28 changes: 28 additions & 0 deletions rocketmq-client/src/consumer/store/offset_serialize_wrapper.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,28 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one or more
* contributor license agreements. See the NOTICE file distributed with
* this work for additional information regarding copyright ownership.
* The ASF licenses this file to You under the Apache License, Version 2.0
* (the "License"); you may not use this file except in compliance with
* the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
use std::collections::HashMap;
use std::sync::atomic::AtomicI64;

use rocketmq_common::common::message::message_queue::MessageQueue;
use serde::Deserialize;
use serde::Serialize;

#[derive(Serialize, Deserialize, Debug, Default)]

Check warning on line 24 in rocketmq-client/src/consumer/store/offset_serialize_wrapper.rs

View check run for this annotation

Codecov / codecov/patch

rocketmq-client/src/consumer/store/offset_serialize_wrapper.rs#L24

Added line #L24 was not covered by tests
#[serde(rename_all = "camelCase")]
pub struct OffsetSerializeWrapper {
pub offset_table: HashMap<MessageQueue, AtomicI64>,

Check warning on line 27 in rocketmq-client/src/consumer/store/offset_serialize_wrapper.rs

View check run for this annotation

Codecov / codecov/patch

rocketmq-client/src/consumer/store/offset_serialize_wrapper.rs#L27

Added line #L27 was not covered by tests
}

0 comments on commit 8c74161

Please sign in to comment.