Merge pull request #45 from CoLearn-Dev/redis-stream

stneng · web-flow · commit 4e6a2cce1159 · 2023-02-10T11:23:33.000+08:00
MQ alternative: redis stream
diff --git a/.github/workflows/check.yml b/.github/workflows/check.yml
@@ -12,13 +12,25 @@ jobs:
     defaults:
       run:
         shell: bash
+    strategy:
+      matrix:
+        mq: [standalone, rabbitmq, redis]
+        include:
+          - mq: rabbitmq
+            docker_image: "rabbitmq:3.8-management"
+            mq_uri: "amqp://guest:guest@localhost"
+            mq_api: "http://guest:guest@localhost:15672/api"
+          - mq: redis
+            docker_image: "redis"
+            mq_uri: "redis://localhost:16379"
     services:
-      rabbitmq:
-        image: rabbitmq:3.8-management
+      mq:
+        image: ${{ matrix.docker_image }}
         ports:
-          - 5672:5672
-          - 15672:15672
-      redis:
+          - 5672:5672   # rabbitmq
+          - 15672:15672 # rabbitmq
+          - 16379:6379  # redis
+      redis:  # for storage macro
         image: redis
         ports:
           - 6379:6379
@@ -44,4 +56,11 @@ jobs:
         run: bash download-server.sh
         working-directory: tests
       - name: Run tests
-        run: cargo test
+        if: ${{ matrix.mq != 'standalone' }}
+        env:
+          COLINK_SERVER_MQ_URI: ${{ matrix.mq_uri }}
+          COLINK_SERVER_MQ_API: ${{ matrix.mq_api }}
+        run: cargo test test_main # remove test_main after updating protocols(policy module, remote storage)
+      - name: Run tests (standalone)
+        if: ${{ matrix.mq == 'standalone' }}
+        run: cargo test test_main # remove test_main after updating protocols(policy module, remote storage)
diff --git a/Cargo.toml b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "colink"
-version = "0.2.11"
+version = "0.3.0"
 edition = "2021"
 description = "CoLink Rust SDK"
 license = "MIT"
@@ -23,7 +23,7 @@ lapin = "2.1"
 prost = "0.10"
 rand = { version = "0.8", features = ["std_rng"] }
 rcgen = { version = "0.10", optional = true }
-redis = { version = "0.22", features = ["tokio-comp"], optional = true }
+redis = { version = "0.22", features = ["tokio-comp"] }
 reqwest = { version = "0.11", default-features = false, features = ["rustls-tls-native-roots"], optional = true }
 secp256k1 = { version = "0.25", features = ["rand-std"] }
 serde = { version = "1.0", features = ["derive"] }
@@ -34,6 +34,7 @@ tokio-rustls = { version = "0.23", optional = true }
 tonic = { version = "0.7", features = ["tls", "tls-roots"] }
 tracing = "0.1"
 tracing-subscriber = "0.2"
+url = "2.2"
 uuid = { version = "0.8", features = ["v4"] }
 
 [build-dependencies]
@@ -48,4 +49,4 @@ variable_transfer = ["extensions", "remote_storage", "hyper", "jsonwebtoken", "r
 registry = []
 policy_module = []
 instant_server = ["reqwest"]
-storage_macro = ["async-recursion", "redis"]
+storage_macro = ["async-recursion"]
diff --git a/README.md b/README.md
@@ -9,7 +9,7 @@ CoLink SDK helps both application and protocol developers access the functionali
 Add this to your Cargo.toml:
 ```toml
 [dependencies]
-colink = "0.2.11"
+colink = "0.3.0"
 ```
 
 ## Getting Started
diff --git a/src/application.rs b/src/application.rs
@@ -4,7 +4,11 @@ use futures_lite::stream::StreamExt;
 use lapin::{
     options::{BasicAckOptions, BasicConsumeOptions},
     types::FieldTable,
-    Connection, ConnectionProperties, Consumer,
+    ConnectionProperties,
+};
+use redis::{
+    streams::{StreamReadOptions, StreamReadReply},
+    AsyncCommands, FromRedisValue,
 };
 use secp256k1::Secp256k1;
 use serde::{Deserialize, Serialize};
@@ -544,32 +548,91 @@ pub struct CoLinkInfo {
     pub version: String,
 }
 
+pub enum CoLinkMQType {
+    RabbitMQ,
+    RedisStream,
+}
 pub struct CoLinkSubscriber {
-    consumer: Consumer,
+    mq_type: CoLinkMQType,
+    queue_name: String,
+    rabbitmq_consumer: Option<lapin::Consumer>,
+    redis_connection: Option<redis::aio::Connection>,
 }
 
 impl CoLinkSubscriber {
     pub async fn new(mq_uri: &str, queue_name: &str) -> Result<Self, Error> {
-        let mq = Connection::connect(mq_uri, ConnectionProperties::default()).await?;
-        let channel = mq.create_channel().await?;
-        let consumer = channel
-            .basic_consume(
-                queue_name,
-                "",
-                BasicConsumeOptions::default(),
-                FieldTable::default(),
-            )
-            .await?;
-        Ok(Self { consumer })
+        let uri_parsed = url::Url::parse(mq_uri)?;
+        if uri_parsed.scheme().starts_with("redis") {
+            let client = redis::Client::open(mq_uri)?;
+            let con = client.get_async_connection().await?;
+            Ok(Self {
+                mq_type: CoLinkMQType::RedisStream,
+                queue_name: queue_name.to_string(),
+                rabbitmq_consumer: None,
+                redis_connection: Some(con),
+            })
+        } else {
+            let mq = lapin::Connection::connect(mq_uri, ConnectionProperties::default()).await?;
+            let channel = mq.create_channel().await?;
+            let consumer = channel
+                .basic_consume(
+                    queue_name,
+                    "",
+                    BasicConsumeOptions::default(),
+                    FieldTable::default(),
+                )
+                .await?;
+            Ok(Self {
+                mq_type: CoLinkMQType::RabbitMQ,
+                queue_name: queue_name.to_string(),
+                rabbitmq_consumer: Some(consumer),
+                redis_connection: None,
+            })
+        }
     }
 
     pub async fn get_next(&mut self) -> Result<Vec<u8>, Error> {
-        let delivery = self.consumer.next().await.expect("error in consumer");
-        let delivery = delivery.expect("error in consumer");
-        let data = String::from_utf8_lossy(&delivery.data);
-        debug!("CoLinkSubscriber Received [{}]", data);
-        delivery.ack(BasicAckOptions::default()).await?;
-        Ok(delivery.data)
+        match self.mq_type {
+            CoLinkMQType::RabbitMQ => {
+                let delivery = self
+                    .rabbitmq_consumer
+                    .as_mut()
+                    .unwrap()
+                    .next()
+                    .await
+                    .expect("error in consumer");
+                let delivery = delivery.expect("error in consumer");
+                delivery.ack(BasicAckOptions::default()).await?;
+                Ok(delivery.data)
+            }
+            CoLinkMQType::RedisStream => {
+                let opts = StreamReadOptions::default()
+                    .group(&self.queue_name, uuid::Uuid::new_v4().to_string())
+                    .block(0)
+                    .count(1);
+                let res: StreamReadReply = self
+                    .redis_connection
+                    .as_mut()
+                    .unwrap()
+                    .xread_options(&[&self.queue_name], &[">"], &opts)
+                    .await?;
+                let id = &res.keys[0].ids[0].id;
+                let data: Vec<u8> = FromRedisValue::from_redis_value(
+                    res.keys[0].ids[0].map.get("payload").unwrap(),
+                )?;
+                self.redis_connection
+                    .as_mut()
+                    .unwrap()
+                    .xack(&self.queue_name, &self.queue_name, &[id])
+                    .await?;
+                self.redis_connection
+                    .as_mut()
+                    .unwrap()
+                    .xdel(&self.queue_name, &[id])
+                    .await?;
+                Ok(data)
+            }
+        }
     }
 }
 
diff --git a/src/extensions/instant_server.rs b/src/extensions/instant_server.rs
@@ -48,7 +48,7 @@ impl InstantServer {
                 .arg("bash -c \"$(curl -fsSL https://raw.githubusercontent.com/CoLearn-Dev/colinkctl/main/install_colink.sh)\"")
                 .env("COLINK_INSTALL_SERVER_ONLY", "true")
                 .env("COLINK_INSTALL_SILENT", "true")
-                .env("COLINK_SERVER_VERSION", "v0.2.9")
+                .env("COLINK_SERVER_VERSION", "v0.3.0")
                 .status()
                 .unwrap();
         }
@@ -61,48 +61,68 @@ impl InstantServer {
             .join("instant_servers")
             .join(instant_server_id.clone());
         std::fs::create_dir_all(&working_dir).unwrap();
-        let mq_amqp = if std::env::var("COLINK_SERVER_MQ_AMQP").is_ok() {
-            std::env::var("COLINK_SERVER_MQ_AMQP").unwrap()
+        let mq_uri = if std::env::var("COLINK_SERVER_MQ_URI").is_ok() {
+            Some(std::env::var("COLINK_SERVER_MQ_URI").unwrap())
         } else {
-            "amqp://guest:guest@localhost:5672".to_string()
+            None
         };
         let mq_api = if std::env::var("COLINK_SERVER_MQ_API").is_ok() {
-            std::env::var("COLINK_SERVER_MQ_API").unwrap()
+            Some(std::env::var("COLINK_SERVER_MQ_API").unwrap())
         } else {
-            "http://guest:guest@localhost:15672/api".to_string()
+            None
         };
-        let (mq_amqp, mq_api) = std::thread::spawn(move || {
+        let (mq_uri, mq_api) = std::thread::spawn(move || {
             tokio::runtime::Builder::new_multi_thread()
                 .enable_all()
                 .build()
                 .unwrap()
                 .block_on(async {
-                    let res = reqwest::get(&mq_api).await.unwrap();
-                    assert!(res.status() == hyper::StatusCode::OK);
-                    lapin::Connection::connect(&mq_amqp, lapin::ConnectionProperties::default())
-                        .await
-                        .unwrap();
+                    if mq_uri.is_some() {
+                        let mq_uri = mq_uri.clone().unwrap();
+                        if mq_uri.starts_with("amqp") {
+                            lapin::Connection::connect(
+                                &mq_uri,
+                                lapin::ConnectionProperties::default(),
+                            )
+                            .await
+                            .unwrap();
+                            if mq_api.is_some() {
+                                let res = reqwest::get(&mq_api.clone().unwrap()).await.unwrap();
+                                assert!(res.status() == hyper::StatusCode::OK);
+                            }
+                        } else if mq_uri.starts_with("redis") {
+                            let client = redis::Client::open(mq_uri).unwrap();
+                            let _con = client.get_async_connection().await.unwrap();
+                        } else {
+                            panic!("mq_uri({}) is not supported.", mq_uri);
+                        }
+                    }
                 });
-            (mq_amqp, mq_api)
+            (mq_uri, mq_api)
         })
         .join()
         .unwrap();
+        let mut args = vec![
+            "--address".to_string(),
+            "0.0.0.0".to_string(),
+            "--port".to_string(),
+            port.to_string(),
+            "--mq-prefix".to_string(),
+            format!("colink-instant-server-{}", port),
+            "--core-uri".to_string(),
+            format!("http://127.0.0.1:{}", port),
+            "--inter-core-reverse-mode".to_string(),
+        ];
+        if let Some(mq_uri) = mq_uri {
+            args.push("--mq-uri".to_string());
+            args.push(mq_uri);
+        }
+        if let Some(mq_api) = mq_api {
+            args.push("--mq-api".to_string());
+            args.push(mq_api);
+        }
         let child = Command::new(program)
-            .args([
-                "--address",
-                "0.0.0.0",
-                "--port",
-                &port.to_string(),
-                "--mq-amqp",
-                &mq_amqp,
-                "--mq-api",
-                &mq_api,
-                "--mq-prefix",
-                &format!("colink-instant-server-{}", port),
-                "--core-uri",
-                &format!("http://127.0.0.1:{}", port),
-                "--inter-core-reverse-mode",
-            ])
+            .args(&args)
             .env("COLINK_HOME", colink_home)
             .current_dir(working_dir.clone())
             .spawn()
diff --git a/src/protocol.rs b/src/protocol.rs
@@ -1,20 +1,14 @@
 use crate::{application::*, utils::get_path_timestamp};
 pub use async_trait::async_trait;
 use clap::Parser;
-use futures_lite::stream::StreamExt;
-use lapin::{
-    options::{BasicAckOptions, BasicConsumeOptions, BasicQosOptions},
-    types::FieldTable,
-    Connection, ConnectionProperties, Consumer,
-};
 use prost::Message;
 use rand::Rng;
 use std::{
     collections::{HashMap, HashSet},
     sync::{Arc, Mutex},
     thread,
 };
-use tracing::{debug, error};
+use tracing::error;
 
 type Error = Box<dyn std::error::Error + Send + Sync + 'static>;
 
@@ -50,12 +44,9 @@ impl CoLinkProtocol {
     }
 
     pub async fn start(&self) -> Result<(), Error> {
-        let mut consumer = self.get_mq_consumer().await?;
-        while let Some(delivery) = consumer.next().await {
-            let delivery = delivery.expect("error in consumer");
-            let data = String::from_utf8_lossy(&delivery.data);
-            debug!("Received [{}]", data);
-            let message: SubscriptionMessage = prost::Message::decode(&*delivery.data).unwrap();
+        let mut subscriber = self.get_subscriber().await?;
+        while let Ok(data) = subscriber.get_next().await {
+            let message: SubscriptionMessage = prost::Message::decode(&*data).unwrap();
             if message.change_type != "delete" {
                 let task_id: Task = prost::Message::decode(&*message.payload).unwrap();
                 let res = self
@@ -74,7 +65,6 @@ impl CoLinkProtocol {
                     Err(e) => error!("Pull Task Error: {}.", e),
                 }
             }
-            delivery.ack(BasicAckOptions::default()).await.unwrap();
         }
 
         Ok(())
@@ -115,7 +105,7 @@ impl CoLinkProtocol {
         Ok(())
     }
 
-    async fn get_mq_consumer(&self) -> Result<Consumer, Error> {
+    async fn get_subscriber(&self) -> Result<CoLinkSubscriber, Error> {
         let operator_mq_key = format!("_internal:protocols:{}:operator_mq", self.protocol_and_role);
         let lock = self.cl.lock(&operator_mq_key).await?;
         let res = self
@@ -172,19 +162,8 @@ impl CoLinkProtocol {
         };
         self.cl.unlock(lock).await?;
 
-        let mq_addr = self.cl.request_info().await?.mq_uri;
-        let mq = Connection::connect(&mq_addr, ConnectionProperties::default()).await?;
-        let channel = mq.create_channel().await?;
-        channel.basic_qos(1, BasicQosOptions::default()).await?;
-        let consumer = channel
-            .basic_consume(
-                &queue_name,
-                "",
-                BasicConsumeOptions::default(),
-                FieldTable::default(),
-            )
-            .await?;
-        Ok(consumer)
+        let subscriber = self.cl.new_subscriber(&queue_name).await?;
+        Ok(subscriber)
     }
 }
 
diff --git a/tests/download-server.sh b/tests/download-server.sh
@@ -2,7 +2,7 @@
 set -e
 rm -rf colink-server
 mkdir colink-server && cd colink-server
-wget https://github.com/CoLearn-Dev/colink-server-dev/releases/download/v0.2.9/colink-server-linux-x86_64.tar.gz
+wget https://github.com/CoLearn-Dev/colink-server-dev/releases/download/v0.3.0/colink-server-linux-x86_64.tar.gz
 tar -xzf colink-server-linux-x86_64.tar.gz
 touch user_init_config.toml # create an empty user init config to prevent automatically starting protocols when importing users.
 cd ..
diff --git a/tests/test_example_protocol_greetings.rs b/tests/test_example_protocol_greetings.rs
