Merge branch 'feature_aqm_codel' into 'rel'

Feature aqm codel

See merge request tango/shaping-engine!76

9 files changed, 273 insertions, 62 deletions

diff --git a/shaping/include/shaper.h b/shaping/include/shaper.h
index 83f93f1..70fded1 100644
--- a/shaping/include/shaper.h
+++ b/shaping/include/shaper.h
@@ -7,6 +7,7 @@
 #include "utils.h"
 #include "shaper_stat.h"
 #include "shaper_global_stat.h"
+#include "shaper_aqm.h"
 extern "C" {
 #include "timeout.h"
 }
@@ -97,18 +98,6 @@ enum shaping_profile_type {
     PROFILE_TYPE_SPLIT_BY_LOCAL_HOST
 };
 
-enum shaper_aqm_type {
-    AQM_TYPE_NONE = 0,
-    AQM_TYPE_BLUE,
-    AQM_TYPE_CODEL,
-    AQM_TYPE_MAX
-};
-
-struct shaper_aqm_blue_para {
-    time_t update_time;
-    int probability;
-};
-
 struct shaper_token_multiple {
     int token_get_multiple;
     unsigned char has_drop_by_queue_full;
@@ -131,6 +120,7 @@ struct shaping_profile_hash_node {
     int tconsume_ref_cnt;
     struct shaper_token_multiple token_multiple;
     struct shaper_aqm_blue_para aqm_blue_para;
+    struct shaper_aqm_codel_para aqm_codel_para;
     unsigned char is_invalid;
     struct timeout timeout_handle;
     UT_hash_handle hh;
diff --git a/shaping/include/shaper_aqm.h b/shaping/include/shaper_aqm.h
index 794c74d..2955c3a 100644
--- a/shaping/include/shaper_aqm.h
+++ b/shaping/include/shaper_aqm.h
@@ -1,2 +1,39 @@
+#pragma once
+#include <time.h>
 
-int shaper_aqm_need_drop(struct shaping_profile_info *profile, struct shaping_packet_wrapper *pkt_wrapper);
\ No newline at end of file
+#define BLUE_PROBABILITY_MAX 100
+#define BLUE_INCREMENT 10
+#define BLUE_DECREMENT 1
+#define BLUE_FREEZE_TIME 2 //unit:s
+#define BLUE_QUEUE_LEN_MAX 100
+
+#define CODEL_MAX_LATENCY 5 //unit:ms
+#define CODEL_DROP_INTERVAL 100 //unit:ms
+enum shaper_aqm_type {
+    AQM_TYPE_NONE = 0,
+    AQM_TYPE_BLUE,
+    AQM_TYPE_CODEL,
+    AQM_TYPE_MAX
+};
+
+struct shaper_aqm_blue_para {
+    time_t update_time;
+    int probability;
+};
+
+enum CoDel_state {
+    CODEL_STATE_NORMAL = 0,
+    CODEL_STATE_DROPPING_TIMER,
+    CODEL_STATE_DROPPING_PHASE,
+};
+
+struct shaper_aqm_codel_para {
+    unsigned long long start_drop_time_ms;
+    unsigned long long next_drop_time_ms;
+    enum CoDel_state state;
+    unsigned int drop_count;
+};
+
+int shaper_aqm_need_drop(struct shaping_profile_info *profile, struct shaping_packet_wrapper *pkt_wrapper, struct timespec *curr_time, unsigned long long latency_us);
+int shaper_aqm_blue_need_drop(struct shaper_aqm_blue_para *para, int curr_queue_len);
+int shaper_aqm_codel_need_drop(struct shaper_aqm_codel_para *para, unsigned long long curr_time_ms, unsigned long long latency_ms);
\ No newline at end of file
diff --git a/shaping/include/shaper_maat.h b/shaping/include/shaper_maat.h
index b570b60..0b50590 100644
--- a/shaping/include/shaper_maat.h
+++ b/shaping/include/shaper_maat.h
@@ -18,6 +18,7 @@ struct shaping_rule {
 struct shaping_profile {
     int id;
     enum shaping_profile_type type;
+    enum shaper_aqm_type aqm_type;
     int in_limit_bandwidth;
     int out_limit_bandwidth;
     int valid;
diff --git a/shaping/src/shaper.cpp b/shaping/src/shaper.cpp
index 32b41cb..218dbb6 100644
--- a/shaping/src/shaper.cpp
+++ b/shaping/src/shaper.cpp
@@ -246,7 +246,7 @@ void shaper_queue_clear(struct shaping_flow *sf, struct shaping_thread_ctx *ctx)
 }
 
 //return success(0) while any avl tree insert success
-static int shaper_flow_push(struct shaping_thread_ctx *ctx, struct shaping_flow *sf, unsigned long long enqueue_time)
+static int shaper_flow_push(struct shaping_thread_ctx *ctx, struct shaping_flow *sf, unsigned long long enqueue_time_us)
 {
     struct shaping_node *s_node = (struct shaping_node*)sf;
     struct shaping_rule_info *s_rule_info = &sf->matched_rule_infos[sf->anchor];
@@ -277,7 +277,7 @@ static int shaper_flow_push(struct shaping_thread_ctx *ctx, struct shaping_flow
     }
 
 END:
-    s_rule_info->primary.enqueue_time_us = enqueue_time;
+    s_rule_info->primary.enqueue_time_us = enqueue_time_us;
     shaper_stat_queueing_pkt_inc(&s_rule_info->primary.stat, pkt_wrapper->direction, ctx->thread_index);
 
     return 0;
@@ -580,14 +580,14 @@ static int shaper_deposit_token_is_enough(struct shaping_profile_info *profile,
     }
 }
 
-static int shaper_token_get_from_profile(struct shaping_thread_ctx *ctx, struct shaping_flow *sf, struct shaping_profile_info *pf_info, int profile_type, int req_token_bits, unsigned char direction, struct timespec *curr_timespec)
+static void shaper_token_get_from_profile(struct shaping_thread_ctx *ctx, struct shaping_flow *sf, struct shaping_profile_info *pf_info, int profile_type, int req_token_bits, unsigned char direction, struct timespec *curr_timespec)
 {
     struct shaping_tconsume_cb_arg *arg = NULL;
     struct shaping_profile_hash_node *pf_hash_node = pf_info->hash_node;
     char key[32] = {0};
 
     if (pf_hash_node->tconsume_ref_cnt > 0) {
-        return SHAPER_TOKEN_GET_FAILED;
+        return;
     }
             
     snprintf(key, sizeof(key), "tsg-shaping-%d-%s", pf_info->id, direction == SHAPING_DIR_OUT ? "outgoing" : "incoming");
@@ -626,7 +626,7 @@ static int shaper_token_get_from_profile(struct shaping_thread_ctx *ctx, struct
             break;
     }
 
-    return SHAPER_TOKEN_GET_FAILED;
+    return;
 }
 
 static void shaper_queue_len_get_cb(const struct swarmkv_reply *reply, void * cb_arg)
@@ -799,9 +799,15 @@ static int shaper_token_consume(struct shaping_thread_ctx *ctx, struct shaping_f
 
     if (shaper_profile_is_priority_blocked(ctx, sf, profile, curr_timespec, curr_time_ms)) {
         return SHAPER_TOKEN_GET_FAILED;
+    }
+
+    int req_token_bits = req_token_bytes * 8;
+    shaper_token_get_from_profile(ctx, sf, profile, profile_type, req_token_bits, direction, curr_timespec);
+
+    if (profile->hash_node->is_invalid && profile_type == PROFILE_IN_RULE_TYPE_PRIMARY) {//for primary, means this rule don't need get token
+        return SHAPER_TOKEN_GET_SUCCESS;
     } else {
-        int req_token_bits = req_token_bytes * 8;
-        return shaper_token_get_from_profile(ctx, sf, profile, profile_type, req_token_bits, direction, curr_timespec);
+        return SHAPER_TOKEN_GET_FAILED;
     }
 }
 
@@ -863,7 +869,8 @@ static enum shaping_packet_action shaper_pkt_action_decide_queueing(struct shapi
     struct shaping_packet_wrapper *pkt_wrapper = NULL;
     struct shaping_profile_container pf_container[SHAPING_PRIORITY_NUM_MAX];
     struct timespec curr_time;
-    unsigned long long enqueue_time;
+    unsigned long long enqueue_time_us;
+    unsigned long long latency_us;
     int get_token_success = 0;
     int profile_num;
 
@@ -875,8 +882,9 @@ static enum shaping_packet_action shaper_pkt_action_decide_queueing(struct shapi
 
     clock_gettime(CLOCK_MONOTONIC, &curr_time);
 
+    latency_us = shaper_pkt_latency_us_calculate(&rule->primary, &curr_time);
     if (pf_container[0].pf_type == PROFILE_IN_RULE_TYPE_PRIMARY) {
-        if (shaper_pkt_latency_us_calculate(pf_container[0].pf_info, &curr_time) > ctx->conf.pkt_max_delay_time_us) {
+        if (latency_us > ctx->conf.pkt_max_delay_time_us) {
             shaper_flow_pop(ctx, sf, &curr_time);
             goto DROP;
         }
@@ -887,7 +895,7 @@ static enum shaping_packet_action shaper_pkt_action_decide_queueing(struct shapi
         profile_type = pf_container[i].pf_type;
 
         /*AQM process, if aqm not pass, for primary profile drop packet, for borrow profile just don't give token to this packet*/
-        if (shaper_aqm_need_drop(profile, pkt_wrapper)) {
+        if (shaper_aqm_need_drop(profile, pkt_wrapper, &curr_time, latency_us)) {
             if (profile_type == PROFILE_IN_RULE_TYPE_PRIMARY) {
                 shaper_flow_pop(ctx, sf, &curr_time);
                 goto DROP;
@@ -918,8 +926,8 @@ static enum shaping_packet_action shaper_pkt_action_decide_queueing(struct shapi
     }
 
     //push sf for next rule
-    enqueue_time = curr_time.tv_sec * MICRO_SECONDS_PER_SEC + curr_time.tv_nsec / NANO_SECONDS_PER_MICRO_SEC;
-    if (0 == shaper_flow_push(ctx, sf, enqueue_time)) {
+    enqueue_time_us = curr_time.tv_sec * MICRO_SECONDS_PER_SEC + curr_time.tv_nsec / NANO_SECONDS_PER_MICRO_SEC;
+    if (0 == shaper_flow_push(ctx, sf, enqueue_time_us)) {
         return SHAPING_QUEUED;
     } else {
         goto DROP;
@@ -936,7 +944,7 @@ static enum shaping_packet_action shaper_pkt_action_decide_no_queue(struct shapi
 {
     int profile_type = PROFILE_IN_RULE_TYPE_PRIMARY;
     struct timespec curr_time;
-    unsigned long long enqueue_time;
+    unsigned long long enqueue_time_us;
     int enqueue_success = 0;
 
     clock_gettime(CLOCK_MONOTONIC, &curr_time);
@@ -964,8 +972,8 @@ static enum shaping_packet_action shaper_pkt_action_decide_no_queue(struct shapi
         goto DROP;
     }
 
-    enqueue_time = curr_time.tv_sec * MICRO_SECONDS_PER_SEC + curr_time.tv_nsec / NANO_SECONDS_PER_MICRO_SEC;
-    if (0 == shaper_flow_push(ctx, sf, enqueue_time)) {
+    enqueue_time_us = curr_time.tv_sec * MICRO_SECONDS_PER_SEC + curr_time.tv_nsec / NANO_SECONDS_PER_MICRO_SEC;
+    if (0 == shaper_flow_push(ctx, sf, enqueue_time_us)) {
         return SHAPING_QUEUED;
     } else {
         goto DROP;
diff --git a/shaping/src/shaper_aqm.cpp b/shaping/src/shaper_aqm.cpp
index 20c4190..e06bf7e 100644
--- a/shaping/src/shaper_aqm.cpp
+++ b/shaping/src/shaper_aqm.cpp
@@ -1,32 +1,65 @@
+#include <math.h>
 #include <time.h>
 #include "shaper.h"
 #include "shaper_aqm.h"
 
-#define PROBABILITY_MAX 100
-#define INCREMENT 10
-#define DECREMENT 1
-#define FREEZE_TIME 2 //unit:s
-#define QUEUE_LEN_MAX 100
-static int shaper_aqm_blue_need_drop(struct shaping_packet_wrapper *pkt_wrapper, struct shaper_aqm_blue_para *para, int curr_queue_len)
+thread_local unsigned int seed = 0;
+int shaper_aqm_blue_need_drop(struct shaper_aqm_blue_para *para, int curr_queue_len)
 {
     time_t curr_time;
 
-    if (time(&curr_time) - para->update_time >= FREEZE_TIME) {
+    if (time(&curr_time) - para->update_time >= BLUE_FREEZE_TIME) {
         para->update_time = curr_time;
-        if (curr_queue_len >= QUEUE_LEN_MAX) {
-            para->probability = (para->probability + INCREMENT) > PROBABILITY_MAX ? PROBABILITY_MAX : (para->probability + INCREMENT);
+        if (curr_queue_len >= BLUE_QUEUE_LEN_MAX) {
+            para->probability = (para->probability + BLUE_INCREMENT) > BLUE_PROBABILITY_MAX ? BLUE_PROBABILITY_MAX : (para->probability + BLUE_INCREMENT);
         } else if (curr_queue_len == 0) {
-            para->probability = (para->probability - DECREMENT) >= 0 ? (para->probability - DECREMENT) : 0;
+            para->probability = (para->probability - BLUE_DECREMENT) >= 0 ? (para->probability - BLUE_DECREMENT) : 0;
         }
     }
 
-    if (rand() % PROBABILITY_MAX < para->probability) {
+    if (rand_r(&seed) % BLUE_PROBABILITY_MAX < para->probability) {
         return 1;
     }
 
     return 0;
 }
 
+int shaper_aqm_codel_need_drop(struct shaper_aqm_codel_para *para, unsigned long long curr_time_ms, unsigned long long latency_ms)
+{
+    if (latency_ms < CODEL_MAX_LATENCY) {
+        para->state = CODEL_STATE_NORMAL;
+        return 0;
+    } 
+
+    int ret = 0;
+    switch (para->state) {
+        case CODEL_STATE_NORMAL:
+            para->start_drop_time_ms = curr_time_ms + CODEL_DROP_INTERVAL;
+            para->state = CODEL_STATE_DROPPING_TIMER;
+            break;
+        case CODEL_STATE_DROPPING_TIMER:
+            if (curr_time_ms >= para->start_drop_time_ms) {
+                para->state = CODEL_STATE_DROPPING_PHASE;
+                para->drop_count = 1;
+                para->next_drop_time_ms = curr_time_ms + CODEL_DROP_INTERVAL / sqrt(para->drop_count);
+                ret = 1;
+            }
+            break;
+        case CODEL_STATE_DROPPING_PHASE:
+            if (curr_time_ms >= para->next_drop_time_ms) {
+                para->drop_count++;
+                para->next_drop_time_ms = curr_time_ms + CODEL_DROP_INTERVAL / sqrt(para->drop_count);
+                ret = 1;
+            }
+            break;
+        default:
+            break;
+    }
+
+    return ret;
+}
+
+
 static int shaper_aqm_have_processed(struct shaping_packet_wrapper *pkt_wrapper, int profile_id)
 {
     int i = 0;
@@ -54,9 +87,10 @@ static void shaper_aqm_mark_processed(struct shaping_packet_wrapper *pkt_wrapper
     }
 }
 
-int shaper_aqm_need_drop(struct shaping_profile_info *profile, struct shaping_packet_wrapper *pkt_wrapper)
+int shaper_aqm_need_drop(struct shaping_profile_info *profile, struct shaping_packet_wrapper *pkt_wrapper, struct timespec *curr_time, unsigned long long latency_us)
 {
     int ret = 0;
+    unsigned long long curr_time_ms;
     
     if (profile->hash_node->aqm_type == AQM_TYPE_NONE) {
         return 0;
@@ -68,9 +102,11 @@ int shaper_aqm_need_drop(struct shaping_profile_info *profile, struct shaping_pa
     
     switch (profile->hash_node->aqm_type) {
         case AQM_TYPE_BLUE:
-            ret = shaper_aqm_blue_need_drop(pkt_wrapper, &profile->hash_node->aqm_blue_para, profile->hash_node->queue_len[profile->priority]);
+            ret = shaper_aqm_blue_need_drop(&profile->hash_node->aqm_blue_para, profile->hash_node->queue_len[profile->priority]);
             break;
         case AQM_TYPE_CODEL:
+            curr_time_ms = curr_time->tv_sec * MILLI_SECONDS_PER_SEC + curr_time->tv_nsec / NANO_SECONDS_PER_MILLI_SEC;
+            ret = shaper_aqm_codel_need_drop(&profile->hash_node->aqm_codel_para, curr_time_ms, latency_us / 1000);
             break;
         default:
             break;
diff --git a/shaping/src/shaper_maat.cpp b/shaping/src/shaper_maat.cpp
index 8a0fa88..ed68844 100644
--- a/shaping/src/shaper_maat.cpp
+++ b/shaping/src/shaper_maat.cpp
@@ -191,6 +191,34 @@ void shaper_profile_ex_new(const char *table_name, int table_id, const char *key
         goto END;
     }
 
+    //parse aqm options
+    json = cJSON_Parse(aqm_options);
+    if (!json) {
+        LOG_ERROR("%s: json parse profile aqm options failed for profile id %d, line %s", LOG_TAG_MAAT, s_pf->id, table_line);
+        goto END;
+    }
+
+    tmp_obj = cJSON_GetObjectItem(json, "enabled");
+    if (!tmp_obj) {
+        LOG_ERROR("%s: json parse aqm enabled failed for profile id %d, line %s", LOG_TAG_MAAT, s_pf->id, table_line);
+        goto END;
+    }
+    if (tmp_obj->valueint == 1) {
+        tmp_obj = cJSON_GetObjectItem(json, "algorithm");
+        if (!tmp_obj) {
+            LOG_ERROR("%s: json parse algorithm failed for profile id %d, line %s", LOG_TAG_MAAT, s_pf->id, table_line);
+            goto END;
+        }
+        if (strncmp(tmp_obj->valuestring, "blue", strlen(tmp_obj->valuestring)) == 0) {
+            s_pf->aqm_type = AQM_TYPE_BLUE;
+        } else if (strncmp(tmp_obj->valuestring, "codel", strlen(tmp_obj->valuestring)) == 0) {
+            s_pf->aqm_type = AQM_TYPE_CODEL;
+        } else {
+            LOG_ERROR("%s: json parse aqm type wrong for profile id %d, line %s", LOG_TAG_MAAT, s_pf->id, table_line);
+            goto END;
+        }
+    }
+    cJSON_Delete(json);
     
     //parse limits of profile
     json = cJSON_Parse(limits);
@@ -258,6 +286,7 @@ void shaper_profile_update(struct shaping_thread_ctx *ctx, struct shaping_profil
     s_pf_info->id = s_pf_ex->id;
     s_pf_info->type = s_pf_ex->type;
     shaper_profile_hash_node_update(ctx, s_pf_info);
+    s_pf_info->hash_node->aqm_type = s_pf_ex->aqm_type;
 
     return;
 }
diff --git a/shaping/test/CMakeLists.txt b/shaping/test/CMakeLists.txt
index 4a9436f..1c75bd6 100644
--- a/shaping/test/CMakeLists.txt
+++ b/shaping/test/CMakeLists.txt
@@ -27,11 +27,21 @@ target_include_directories(gtest_shaper PUBLIC ${CMAKE_SOURCE_DIR}/shaping/inclu
 target_link_libraries(gtest_shaper common shaper pthread gtest)
 
 ###############################################################################
+# gtest_shaper_aqm
+###############################################################################
+add_executable(gtest_shaper_aqm gtest_shaper_aqm.cpp)
+target_include_directories(gtest_shaper_aqm PUBLIC ${CMAKE_SOURCE_DIR}/common/include)
+target_include_directories(gtest_shaper_aqm PUBLIC ${CMAKE_SOURCE_DIR}/shaping/include)
+target_link_libraries(gtest_shaper_aqm common shaper pthread gtest)
+
+###############################################################################
 # gtest_discover_tests
 ###############################################################################
 
 include(GoogleTest)
 gtest_discover_tests(gtest_shaper_maat)
+gtest_discover_tests(gtest_shaper_send_log)
 gtest_discover_tests(gtest_shaper)
+gtest_discover_tests(gtest_shaper_aqm)
 
 file(COPY ./test_conf/ DESTINATION ./conf/)
\ No newline at end of file
diff --git a/shaping/test/gtest_shaper_aqm.cpp b/shaping/test/gtest_shaper_aqm.cpp
new file mode 100644
index 0000000..92bab10
--- /dev/null
+++ b/shaping/test/gtest_shaper_aqm.cpp
@@ -0,0 +1,96 @@
+#include <gtest/gtest.h>
+#include <unistd.h>
+
+#include "shaper.h"
+#include "shaper_aqm.h"
+
+TEST(aqm_bule, need_drop)
+{
+    struct shaper_aqm_blue_para para;
+    int drop_cnt = 0;
+
+    para.update_time = 0;
+    para.probability = 0;
+
+    for (int i = 0; i < 10000; i++) {
+        if (shaper_aqm_blue_need_drop(&para, BLUE_QUEUE_LEN_MAX + 1)) {
+            drop_cnt++;
+        }
+    }
+    
+    EXPECT_EQ(para.probability, BLUE_INCREMENT);
+
+    EXPECT_GT(drop_cnt, 900);
+    EXPECT_LT(drop_cnt, 1100);
+
+    sleep(3);
+    shaper_aqm_blue_need_drop(&para, 0);
+    EXPECT_EQ(para.probability, BLUE_INCREMENT - BLUE_DECREMENT);
+}
+
+TEST(aqm_blue, no_drop)
+{
+    struct shaper_aqm_blue_para para;
+    int drop_cnt = 0;
+
+    para.update_time = 0;
+    para.probability = 0;
+
+    for (int i = 0; i < 10000; i++) {
+        if (shaper_aqm_blue_need_drop(&para, BLUE_QUEUE_LEN_MAX - 1)) {
+            drop_cnt++;
+        }
+    }
+    
+    EXPECT_EQ(para.probability, 0);
+    EXPECT_EQ(drop_cnt, 0);
+
+    sleep(3);
+    shaper_aqm_blue_need_drop(&para, BLUE_QUEUE_LEN_MAX - 1);
+    EXPECT_EQ(para.probability, 0);
+}
+
+TEST(aqm_codel, need_drop)
+{
+    struct shaper_aqm_codel_para para;
+    int curr_time_ms = 0;
+
+    memset(&para, 0, sizeof(para));
+
+    shaper_aqm_codel_need_drop(&para, curr_time_ms, CODEL_MAX_LATENCY + 1);
+    EXPECT_EQ(para.state, CODEL_STATE_DROPPING_TIMER);
+    EXPECT_EQ(para.start_drop_time_ms, curr_time_ms + CODEL_DROP_INTERVAL);
+
+    curr_time_ms = para.start_drop_time_ms + 1;
+    shaper_aqm_codel_need_drop(&para, curr_time_ms, CODEL_MAX_LATENCY + 1);
+    EXPECT_EQ(para.state, CODEL_STATE_DROPPING_PHASE);
+    EXPECT_EQ(para.drop_count, 1);
+    EXPECT_EQ(para.next_drop_time_ms, int(curr_time_ms + CODEL_DROP_INTERVAL / sqrt(para.drop_count)));
+
+    curr_time_ms = para.next_drop_time_ms + 1;
+    shaper_aqm_codel_need_drop(&para, curr_time_ms, CODEL_MAX_LATENCY + 1);
+    EXPECT_EQ(para.state, CODEL_STATE_DROPPING_PHASE);
+    EXPECT_EQ(para.drop_count, 2);
+    EXPECT_EQ(para.next_drop_time_ms, int(curr_time_ms + CODEL_DROP_INTERVAL / sqrt(para.drop_count)));
+
+    shaper_aqm_codel_need_drop(&para, curr_time_ms, CODEL_MAX_LATENCY - 1);
+    EXPECT_EQ(para.state, CODEL_STATE_NORMAL);
+}
+
+TEST(aqm_codel, no_drop)
+{
+    struct shaper_aqm_codel_para para;
+    int curr_time_ms = 0;
+
+    memset(&para, 0, sizeof(para));
+
+    shaper_aqm_codel_need_drop(&para, curr_time_ms, CODEL_MAX_LATENCY - 1);
+    EXPECT_EQ(para.state, CODEL_STATE_NORMAL);
+    EXPECT_EQ(para.drop_count, 0);
+}
+
+int main(int argc, char **argv)
+{
+	::testing::InitGoogleTest(&argc, argv);
+	return RUN_ALL_TESTS();
+}
\ No newline at end of file
diff --git a/shaping/test/gtest_shaper_maat.cpp b/shaping/test/gtest_shaper_maat.cpp
index ff94efe..770c937 100644
--- a/shaping/test/gtest_shaper_maat.cpp
+++ b/shaping/test/gtest_shaper_maat.cpp
@@ -54,6 +54,7 @@ TEST(shaping_profile, parse)
     EXPECT_EQ(s_pf->in_limit_bandwidth, 1024);
     EXPECT_EQ(s_pf->out_limit_bandwidth, 2048);
     EXPECT_EQ(s_pf->type, PROFILE_TYPE_MAX_MIN_HOST_FAIRNESS);
+    EXPECT_EQ(s_pf->aqm_type, AQM_TYPE_CODEL);
 
     shaper_profile_ex_dup(0, (void**)&s_pf_dup, (void**)&s_pf, 0, NULL);
     EXPECT_EQ(s_pf, s_pf_dup);
@@ -64,7 +65,7 @@ TEST(shaping_profile, parse)
 
 TEST(shaping_flow, update_rule)
 {
-    struct shaping_thread_ctx ctx;
+    struct shaping_ctx *ctx = NULL;
     struct shaping_flow sf;
     struct shaping_rule_info *rule_info;
     long long rule_ids[] = {1, 2, 3};
@@ -74,15 +75,16 @@ TEST(shaping_flow, update_rule)
 
     stub_init();
 
+    ctx = shaping_engine_init();
+
     stub_set_matched_shaping_rules(3, rule_ids, prioritys, profile_nums, profile_ids);
 
-    ctx.maat_info = (struct shaping_maat_info *)calloc(1, sizeof(struct shaping_maat_info));
-    ctx.maat_info->rule_table_id = STUB_MAAT_SHAPING_RULE_TABLE_ID;
-    ctx.maat_info->profile_table_id = STUB_MAAT_SHAPING_PROFILE_TABLE_ID;
+    ctx->maat_info->rule_table_id = STUB_MAAT_SHAPING_RULE_TABLE_ID;
+    ctx->maat_info->profile_table_id = STUB_MAAT_SHAPING_PROFILE_TABLE_ID;
 
     memset(&sf, 0, sizeof(sf));
     sf.priority = SHAPING_PRIORITY_NUM_MAX;
-    shaper_rules_update(&ctx, &sf, rule_ids, 3);
+    shaper_rules_update(&ctx->thread_ctx[0], &sf, rule_ids, 3);
 
     EXPECT_EQ(sf.rule_num, 3);
 
@@ -110,12 +112,12 @@ TEST(shaping_flow, update_rule)
     EXPECT_EQ(rule_info->borrowing[1].id, 6);
     EXPECT_EQ(rule_info->borrowing[1].priority, 3);
 
-    free(ctx.maat_info);
+    shaping_engine_destroy(ctx);
 }
 
 TEST(shaping_flow, update_rule_dup)
 {
-    struct shaping_thread_ctx ctx;
+    struct shaping_ctx *ctx = NULL;
     struct shaping_flow sf;
     struct shaping_rule_info *rule_info;
     long long rule_ids[] = {1, 2, 3, 4};
@@ -125,17 +127,18 @@ TEST(shaping_flow, update_rule_dup)
 
     stub_init();
 
+    ctx = shaping_engine_init();
+
     stub_set_matched_shaping_rules(4, rule_ids, prioritys, profile_nums, profile_ids);
 
-    ctx.maat_info = (struct shaping_maat_info *)calloc(1, sizeof(struct shaping_maat_info));
-    ctx.maat_info->rule_table_id = STUB_MAAT_SHAPING_RULE_TABLE_ID;
-    ctx.maat_info->profile_table_id = STUB_MAAT_SHAPING_PROFILE_TABLE_ID;
+    ctx->maat_info->rule_table_id = STUB_MAAT_SHAPING_RULE_TABLE_ID;
+    ctx->maat_info->profile_table_id = STUB_MAAT_SHAPING_PROFILE_TABLE_ID;
 
     memset(&sf, 0, sizeof(sf));
     sf.priority = SHAPING_PRIORITY_NUM_MAX;
 
     long long rule_id1[] = {1, 2, 3};
-    shaper_rules_update(&ctx, &sf, rule_id1, 3);
+    shaper_rules_update(&ctx->thread_ctx[0], &sf, rule_id1, 3);
 
     EXPECT_EQ(sf.rule_num, 3);
 
@@ -164,23 +167,23 @@ TEST(shaping_flow, update_rule_dup)
     EXPECT_EQ(rule_info->borrowing[1].priority, 3);
 
     long long rule_id2[] = {1};
-    shaper_rules_update(&ctx, &sf, rule_id2, 1);
+    shaper_rules_update(&ctx->thread_ctx[0], &sf, rule_id2, 1);
     EXPECT_EQ(sf.rule_num, 3);
 
     long long rule_id3[] = {2, 3, 4};
-    shaper_rules_update(&ctx, &sf, rule_id3, 3);
+    shaper_rules_update(&ctx->thread_ctx[0], &sf, rule_id3, 3);
     EXPECT_EQ(sf.rule_num, 4);
     rule_info = &sf.matched_rule_infos[3];
     EXPECT_EQ(rule_info->id, 4);
     EXPECT_EQ(rule_info->primary.id, 7);
     EXPECT_EQ(rule_info->primary.priority, 1);
 
-    free(ctx.maat_info);
+    shaping_engine_destroy(ctx);
 }
 
 TEST(shaping_flow, update_rule_after_priority_confirmed)
 {
-    struct shaping_thread_ctx ctx;
+    struct shaping_ctx *ctx = NULL;
     struct shaping_flow sf;
     struct shaping_rule_info *rule_info;
     long long rule_ids[] = {1, 2, 3};
@@ -190,20 +193,21 @@ TEST(shaping_flow, update_rule_after_priority_confirmed)
 
     stub_init();
 
+    ctx = shaping_engine_init();
+
     stub_set_matched_shaping_rules(3, rule_ids, prioritys, profile_nums, profile_ids);
 
-    ctx.maat_info = (struct shaping_maat_info *)calloc(1, sizeof(struct shaping_maat_info));
-    ctx.maat_info->rule_table_id = STUB_MAAT_SHAPING_RULE_TABLE_ID;
-    ctx.maat_info->profile_table_id = STUB_MAAT_SHAPING_PROFILE_TABLE_ID;
+    ctx->maat_info->rule_table_id = STUB_MAAT_SHAPING_RULE_TABLE_ID;
+    ctx->maat_info->profile_table_id = STUB_MAAT_SHAPING_PROFILE_TABLE_ID;
 
     memset(&sf, 0, sizeof(sf));
     sf.priority = SHAPING_PRIORITY_NUM_MAX;
     long long first_rule_ids[] = {2, 3};
-    shaper_rules_update(&ctx, &sf, first_rule_ids, 2);
+    shaper_rules_update(&ctx->thread_ctx[0], &sf, first_rule_ids, 2);
 
     sf.processed_pkts = CONFIRM_PRIORITY_PKTS + 1;
     long long after_confirm_priority_rule_ids[] = {1};
-    shaper_rules_update(&ctx, &sf, after_confirm_priority_rule_ids, 1);
+    shaper_rules_update(&ctx->thread_ctx[0], &sf, after_confirm_priority_rule_ids, 1);
 
     EXPECT_EQ(sf.rule_num, 3);
     EXPECT_EQ(sf.priority, 2);
@@ -232,7 +236,7 @@ TEST(shaping_flow, update_rule_after_priority_confirmed)
     EXPECT_EQ(rule_info->primary.priority, 2);
     EXPECT_EQ(rule_info->borrowing_num, 0);
 
-    free(ctx.maat_info);
+    shaping_engine_destroy(ctx);
 }
 
 int main(int argc, char **argv)