add bidirectional limit direction

7 files changed, 162 insertions, 77 deletions

diff --git a/shaping/include/shaper.h b/shaping/include/shaper.h
index 70fded1..215f6c4 100644
--- a/shaping/include/shaper.h
+++ b/shaping/include/shaper.h
@@ -98,6 +98,11 @@ enum shaping_profile_type {
     PROFILE_TYPE_SPLIT_BY_LOCAL_HOST
 };
 
+enum shaping_profile_limit_direction {
+    PROFILE_LIMIT_DIRECTION_BIDIRECTION,
+    PROFILE_LIMIT_DIRECTION_INCOMING_OUTGOING,
+};
+
 struct shaper_token_multiple {
     int token_get_multiple;
     unsigned char has_drop_by_queue_full;
@@ -108,8 +113,10 @@ struct shaper_token_multiple {
 struct shaping_profile_hash_node {
     int id;
     enum shaper_aqm_type aqm_type;
+    enum shaping_profile_limit_direction limit_direction;
     long long in_deposit_token_bits[SHAPING_PRIORITY_NUM_MAX];
     long long out_deposit_token_bits[SHAPING_PRIORITY_NUM_MAX];
+    long long bidirection_deposit_token_bits[SHAPING_PRIORITY_NUM_MAX];
     long long last_failed_get_token_ms;
     long long last_hmget_ms;
     long long queue_len[SHAPING_PRIORITY_NUM_MAX];
@@ -132,6 +139,7 @@ struct shaping_profile_info {
     int priority;
     long long in_deposit_token_bits;
     long long out_deposit_token_bits;
+    long long bidirection_deposit_token_bits;
     long long last_failed_get_token_ms;
     unsigned long long enqueue_time_us;//to calculate max latency
     struct shaping_stat_for_profile stat;
diff --git a/shaping/include/shaper_maat.h b/shaping/include/shaper_maat.h
index 0b50590..1c8cb62 100644
--- a/shaping/include/shaper_maat.h
+++ b/shaping/include/shaper_maat.h
@@ -19,8 +19,10 @@ struct shaping_profile {
     int id;
     enum shaping_profile_type type;
     enum shaper_aqm_type aqm_type;
+    enum shaping_profile_limit_direction limit_direction;
     int in_limit_bandwidth;
     int out_limit_bandwidth;
+    int bidirection_limit_bandwidth;
     int valid;
 };
 
diff --git a/shaping/src/shaper.cpp b/shaping/src/shaper.cpp
index 218dbb6..cb4c382 100644
--- a/shaping/src/shaper.cpp
+++ b/shaping/src/shaper.cpp
@@ -383,7 +383,9 @@ static void shaper_deposit_token_add(struct shaping_profile_info *profile, int r
 
     switch (profile->type) {
         case PROFILE_TYPE_GENERIC:
-            if (direction == SHAPING_DIR_IN) {
+            if (pf_hash_node->limit_direction == PROFILE_LIMIT_DIRECTION_BIDIRECTION) {
+                deposit_token = &pf_hash_node->bidirection_deposit_token_bits[priority];
+            } else if (direction == SHAPING_DIR_IN) {
                 deposit_token = &pf_hash_node->in_deposit_token_bits[priority];
             } else {
                 deposit_token = &pf_hash_node->out_deposit_token_bits[priority];
@@ -392,7 +394,9 @@ static void shaper_deposit_token_add(struct shaping_profile_info *profile, int r
         case PROFILE_TYPE_HOST_FARINESS:
         case PROFILE_TYPE_MAX_MIN_HOST_FAIRNESS:
         case PROFILE_TYPE_SPLIT_BY_LOCAL_HOST:
-            if (direction == SHAPING_DIR_IN) {
+            if (pf_hash_node->limit_direction == PROFILE_LIMIT_DIRECTION_BIDIRECTION) {
+                deposit_token = &profile->bidirection_deposit_token_bits;
+            } else if (direction == SHAPING_DIR_IN) {
                 deposit_token = &profile->in_deposit_token_bits;
             } else {
                 deposit_token = &profile->out_deposit_token_bits;
@@ -514,14 +518,16 @@ END:
     return;
 }
 
-static void shaper_deposit_token_sub(struct shaping_profile_info *profile, int req_token_bits, unsigned char direction, int priority)
+static int shaper_deposit_token_get(struct shaping_profile_info *profile, int req_token_bits, unsigned char direction, int priority, int force)
 {
     long long *deposit_token;
     struct shaping_profile_hash_node *pf_hash_node = profile->hash_node;
 
     switch (profile->type) {
         case PROFILE_TYPE_GENERIC:
-            if (direction == SHAPING_DIR_IN) {
+            if (pf_hash_node->limit_direction == PROFILE_LIMIT_DIRECTION_BIDIRECTION) {
+                deposit_token = &pf_hash_node->bidirection_deposit_token_bits[priority];
+            } else if (direction == SHAPING_DIR_IN) {
                 deposit_token = &pf_hash_node->in_deposit_token_bits[priority];
             } else {
                 deposit_token = &pf_hash_node->out_deposit_token_bits[priority];
@@ -530,7 +536,9 @@ static void shaper_deposit_token_sub(struct shaping_profile_info *profile, int r
         case PROFILE_TYPE_HOST_FARINESS:
         case PROFILE_TYPE_MAX_MIN_HOST_FAIRNESS:
         case PROFILE_TYPE_SPLIT_BY_LOCAL_HOST:
-            if (direction == SHAPING_DIR_IN) {
+            if (pf_hash_node->limit_direction == PROFILE_LIMIT_DIRECTION_BIDIRECTION) {
+                deposit_token = &profile->bidirection_deposit_token_bits;
+            } else if (direction == SHAPING_DIR_IN) {
                 deposit_token = &profile->in_deposit_token_bits;
             } else {
                 deposit_token = &profile->out_deposit_token_bits;
@@ -538,46 +546,20 @@ static void shaper_deposit_token_sub(struct shaping_profile_info *profile, int r
             break;
         default:
             LOG_ERROR("%s: invalid profile type %d, profile id %d", LOG_TAG_SHAPING, profile->type, profile->id);
-            return;
+            return 0;
     }
 
-    *deposit_token -= req_token_bits;
-
-    return;
-}
-
-static int shaper_deposit_token_is_enough(struct shaping_profile_info *profile, int req_token_bits, unsigned char direction, int priority)
-{
-    int deposit_token;
-    struct shaping_profile_hash_node *pf_hash_node = profile->hash_node;
-
-    switch (profile->type) {
-        case PROFILE_TYPE_GENERIC:
-            if (direction == SHAPING_DIR_IN) {
-                deposit_token = pf_hash_node->in_deposit_token_bits[priority];
-            } else {
-                deposit_token = pf_hash_node->out_deposit_token_bits[priority];
-            }
-            break;
-        case PROFILE_TYPE_HOST_FARINESS:
-        case PROFILE_TYPE_MAX_MIN_HOST_FAIRNESS:
-        case PROFILE_TYPE_SPLIT_BY_LOCAL_HOST:
-            if (direction == SHAPING_DIR_IN) {
-                deposit_token = profile->in_deposit_token_bits;
-            } else {
-                deposit_token = profile->out_deposit_token_bits;
-            }
-            break;
-        default:
-            LOG_ERROR("%s: invalid profile type %d, profile id %d", LOG_TAG_SHAPING, profile->type, profile->id);
-            return 0;
+    if (force) {
+        *deposit_token -= req_token_bits;
+        return 0;
     }
 
-    if (deposit_token >= req_token_bits) {
-        return 1;
-    } else {
+    if (*deposit_token >= req_token_bits) {
+        *deposit_token -= req_token_bits;
         return 0;
     }
+
+    return -1;
 }
 
 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)
@@ -589,8 +571,12 @@ static void shaper_token_get_from_profile(struct shaping_thread_ctx *ctx, struct
     if (pf_hash_node->tconsume_ref_cnt > 0) {
         return;
     }
-            
-    snprintf(key, sizeof(key), "tsg-shaping-%d-%s", pf_info->id, direction == SHAPING_DIR_OUT ? "outgoing" : "incoming");
+    
+    if (pf_hash_node->limit_direction == PROFILE_LIMIT_DIRECTION_BIDIRECTION) {
+        snprintf(key, sizeof(key), "tsg-shaping-%d-bidirectional", pf_info->id);
+    } else {
+        snprintf(key, sizeof(key), "tsg-shaping-%d-%s", pf_info->id, direction == SHAPING_DIR_OUT ? "outgoing" : "incoming");
+    }
 
     arg = (struct shaping_tconsume_cb_arg *)calloc(1, sizeof(struct shaping_tconsume_cb_arg));
     arg->ctx = ctx;
@@ -786,8 +772,7 @@ static int shaper_token_consume(struct shaping_thread_ctx *ctx, struct shaping_f
         return SHAPER_TOKEN_GET_PASS;//rule is disabled, don't need to get token and forward packet
     }
 
-    if (shaper_deposit_token_is_enough(profile, req_token_bytes * 8, direction, profile->priority)) {
-        shaper_deposit_token_sub(profile, req_token_bytes * 8, direction, profile->priority);
+    if (shaper_deposit_token_get(profile, req_token_bytes * 8, direction, profile->priority, 0) == 0) {
         return SHAPER_TOKEN_GET_SUCCESS;
     }
 
@@ -1063,7 +1048,7 @@ static void shaper_token_consume_force(struct shaping_flow *sf, struct metadata
 
     for (int i = 0; i < sf->rule_num; i++) {
         rule = &sf->matched_rule_infos[i];
-        shaper_deposit_token_sub(&rule->primary, meta->raw_len * 8, meta->dir, rule->primary.priority);
+        shaper_deposit_token_get(&rule->primary, meta->raw_len * 8, meta->dir, rule->primary.priority, 1);
     }
 
     return;
diff --git a/shaping/src/shaper_maat.cpp b/shaping/src/shaper_maat.cpp
index ed68844..54d59c5 100644
--- a/shaping/src/shaper_maat.cpp
+++ b/shaping/src/shaper_maat.cpp
@@ -242,6 +242,14 @@ void shaper_profile_ex_new(const char *table_name, int table_id, const char *key
             LOG_ERROR("%s: json parse limit direction failed for profile id %d, line %s", LOG_TAG_MAAT, s_pf->id, table_line);
             goto END;
         }
+
+        if (strncmp(tmp_obj->valuestring, "bidirectional", strlen(tmp_obj->valuestring)) == 0) {
+            s_pf->limit_direction = PROFILE_LIMIT_DIRECTION_BIDIRECTION;
+            s_pf->bidirection_limit_bandwidth = limit_bandwidth;
+            break;//config either bidirectional or incoming/outgoing limit
+        } 
+        
+        s_pf->limit_direction = PROFILE_LIMIT_DIRECTION_INCOMING_OUTGOING;
         if (strncmp(tmp_obj->valuestring, "incoming", strlen(tmp_obj->valuestring)) == 0) {
             s_pf->in_limit_bandwidth = limit_bandwidth;
         } else {
@@ -287,6 +295,7 @@ void shaper_profile_update(struct shaping_thread_ctx *ctx, struct shaping_profil
     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;
+    s_pf_info->hash_node->limit_direction = s_pf_ex->limit_direction;
 
     return;
 }
diff --git a/shaping/test/gtest_shaper.cpp b/shaping/test/gtest_shaper.cpp
index 2b8f296..fe322b6 100644
--- a/shaping/test/gtest_shaper.cpp
+++ b/shaping/test/gtest_shaper.cpp
@@ -244,7 +244,7 @@ TEST(single_session, udp_tx_in_order)
     ASSERT_TRUE(sf != NULL);
 
     stub_set_matched_shaping_rules(1, rule_id, priority, profile_num, profile_id);
-    stub_set_token_bucket_avl_per_sec(0, 1000, SHAPING_DIR_OUT);
+    stub_set_token_bucket_avl_per_sec(0, 1000, SHAPING_DIR_OUT, PROFILE_LIMIT_DIRECTION_INCOMING_OUTGOING);
     actual_tx_queue = stub_get_tx_queue();
     shaper_rules_update(&ctx->thread_ctx[0], sf, rule_id, 1);
 
@@ -302,6 +302,69 @@ TEST(single_session, udp_tx_in_order)
 
 /*session1 match rule1
   rule1:
+      profile: bidirectional limit 1000*/
+TEST(bidirectional, udp_tx_in_order)
+{
+    struct stub_pkt_queue expec_tx_queue;
+    struct stub_pkt_queue *actual_tx_queue;
+    struct shaping_ctx *ctx = NULL;
+    struct shaping_flow *sf = NULL;
+    long long rule_id[] = {0};
+    int priority[] = {1};
+    int profile_num[] = {1};
+    int profile_id[][MAX_REF_PROFILE] = {{0}};
+
+    TAILQ_INIT(&expec_tx_queue);
+    stub_init();
+    ctx = shaping_engine_init();
+    ASSERT_TRUE(ctx != NULL);
+    sf = shaping_flow_new(&ctx->thread_ctx[0]);
+    ASSERT_TRUE(sf != NULL);
+
+    stub_set_matched_shaping_rules(1, rule_id, priority, profile_num, profile_id);
+    stub_set_token_bucket_avl_per_sec(0, 1000, SHAPING_DIR_OUT, PROFILE_LIMIT_DIRECTION_BIDIRECTION);
+    actual_tx_queue = stub_get_tx_queue();
+    shaper_rules_update(&ctx->thread_ctx[0], sf, rule_id, 1);
+
+    /*******send packets***********/
+    send_packets(&ctx->thread_ctx[0], sf, 10, 100, SHAPING_DIR_OUT, &expec_tx_queue, 1, 0);
+    send_packets(&ctx->thread_ctx[0], sf, 10, 100, SHAPING_DIR_IN, &expec_tx_queue, 1, 0);
+    send_packets(&ctx->thread_ctx[0], sf, 10, 100, SHAPING_DIR_OUT, &expec_tx_queue, 1, 0);
+
+
+    //first 10 out packets
+    ASSERT_EQ(0, judge_packet_eq(&expec_tx_queue, actual_tx_queue, 10));
+    ASSERT_TRUE(TAILQ_EMPTY(actual_tx_queue));
+
+    stub_refresh_token_bucket(0);
+    for (int i = 0; i < 11; i++) {//first polling just request token and don't send pkt
+        polling_entry(ctx->thread_ctx[0].sp, ctx->thread_ctx[0].stat, &ctx->thread_ctx[0]);
+        stub_curr_time_ns_inc(STUB_TIME_INC_FOR_PACKET);
+    }
+    //10 out packets
+    ASSERT_EQ(0, judge_packet_eq(&expec_tx_queue, actual_tx_queue, 10));
+    ASSERT_TRUE(TAILQ_EMPTY(actual_tx_queue));
+
+    stub_refresh_token_bucket(0);
+    for (int i = 0; i < 11; i++) {//first polling just request token and don't send pkt
+        polling_entry(ctx->thread_ctx[0].sp, ctx->thread_ctx[0].stat, &ctx->thread_ctx[0]);
+        stub_curr_time_ns_inc(STUB_TIME_INC_FOR_PACKET);
+    }
+    ASSERT_EQ(0, judge_packet_eq(&expec_tx_queue, actual_tx_queue, 10));
+    ASSERT_TRUE(TAILQ_EMPTY(actual_tx_queue));
+    ASSERT_TRUE(TAILQ_EMPTY(&expec_tx_queue));
+
+    shaping_flow_free(&ctx->thread_ctx[0], sf);
+    fieldstat_global_disable_prometheus_endpoint();
+
+    shaper_thread_resource_clear();
+    shaping_engine_destroy(ctx);
+    stub_clear_matched_shaping_rules();
+
+}
+
+/*session1 match rule1
+  rule1:
       profile: limit 1000*/
 TEST(max_min_host_fairness_profile, udp_tx_in_order)
 {
@@ -322,7 +385,7 @@ TEST(max_min_host_fairness_profile, udp_tx_in_order)
     ASSERT_TRUE(sf != NULL);
 
     stub_set_matched_shaping_rules(1, rule_id, priority, profile_num, profile_id);
-    stub_set_token_bucket_avl_per_sec(0, 1000, SHAPING_DIR_OUT);
+    stub_set_token_bucket_avl_per_sec(0, 1000, SHAPING_DIR_OUT, PROFILE_LIMIT_DIRECTION_INCOMING_OUTGOING);
     stub_set_profile_type(0, PROFILE_TYPE_MAX_MIN_HOST_FAIRNESS);
     actual_tx_queue = stub_get_tx_queue();
     shaper_rules_update(&ctx->thread_ctx[0], sf, rule_id, 1);
@@ -403,7 +466,7 @@ TEST(single_session, tcp_tx_in_order)
     ASSERT_TRUE(sf != NULL);
 
     stub_set_matched_shaping_rules(1, rule_id, priority, profile_num, profile_id);
-    stub_set_token_bucket_avl_per_sec(0, 1000, SHAPING_DIR_OUT);
+    stub_set_token_bucket_avl_per_sec(0, 1000, SHAPING_DIR_OUT, PROFILE_LIMIT_DIRECTION_INCOMING_OUTGOING);
     actual_tx_queue = stub_get_tx_queue();
     shaper_rules_update(&ctx->thread_ctx[0], sf, rule_id, 1);
 
@@ -489,7 +552,7 @@ TEST(single_session, udp_diff_direction)
     ASSERT_TRUE(sf != NULL);
 
     stub_set_matched_shaping_rules(1, rule_id, priority, profile_num, profile_id);
-    stub_set_token_bucket_avl_per_sec(0, 1000, SHAPING_DIR_OUT);
+    stub_set_token_bucket_avl_per_sec(0, 1000, SHAPING_DIR_OUT, PROFILE_LIMIT_DIRECTION_INCOMING_OUTGOING);
     actual_tx_queue = stub_get_tx_queue();
     shaper_rules_update(&ctx->thread_ctx[0], sf, rule_id, 1);
 
@@ -568,9 +631,9 @@ TEST(single_session, udp_multi_rules)
     ASSERT_TRUE(sf != NULL);
 
     stub_set_matched_shaping_rules(3, rule_id, priority, profile_num, profile_id);
-    stub_set_token_bucket_avl_per_sec(0, 3000, SHAPING_DIR_OUT);
-    stub_set_token_bucket_avl_per_sec(1, 2000, SHAPING_DIR_OUT);
-    stub_set_token_bucket_avl_per_sec(2, 1000, SHAPING_DIR_OUT);
+    stub_set_token_bucket_avl_per_sec(0, 3000, SHAPING_DIR_OUT, PROFILE_LIMIT_DIRECTION_INCOMING_OUTGOING);
+    stub_set_token_bucket_avl_per_sec(1, 2000, SHAPING_DIR_OUT, PROFILE_LIMIT_DIRECTION_INCOMING_OUTGOING);
+    stub_set_token_bucket_avl_per_sec(2, 1000, SHAPING_DIR_OUT, PROFILE_LIMIT_DIRECTION_INCOMING_OUTGOING);
     actual_tx_queue = stub_get_tx_queue();
     shaper_rules_update(&ctx->thread_ctx[0], sf, rule_id, 3);
 
@@ -648,8 +711,8 @@ TEST(single_session, udp_borrow)
     ASSERT_TRUE(sf != NULL);
 
     stub_set_matched_shaping_rules(1, rule_id, priority, profile_num, profile_id);
-    stub_set_token_bucket_avl_per_sec(1, 0, SHAPING_DIR_OUT);
-    stub_set_token_bucket_avl_per_sec(2, 1000, SHAPING_DIR_OUT);
+    stub_set_token_bucket_avl_per_sec(1, 0, SHAPING_DIR_OUT, PROFILE_LIMIT_DIRECTION_INCOMING_OUTGOING);
+    stub_set_token_bucket_avl_per_sec(2, 1000, SHAPING_DIR_OUT, PROFILE_LIMIT_DIRECTION_INCOMING_OUTGOING);
     actual_tx_queue = stub_get_tx_queue();
     shaper_rules_update(&ctx->thread_ctx[0], sf, rule_id, 1);
 
@@ -722,9 +785,9 @@ TEST(single_session, udp_borrow_same_priority_9)
     ASSERT_TRUE(sf != NULL);
 
     stub_set_matched_shaping_rules(1, rule_id, priority, profile_num, profile_id);
-    stub_set_token_bucket_avl_per_sec(1, 0, SHAPING_DIR_OUT);
-    stub_set_token_bucket_avl_per_sec(2, 0, SHAPING_DIR_OUT);
-    stub_set_token_bucket_avl_per_sec(3, 1000, SHAPING_DIR_OUT);
+    stub_set_token_bucket_avl_per_sec(1, 0, SHAPING_DIR_OUT, PROFILE_LIMIT_DIRECTION_INCOMING_OUTGOING);
+    stub_set_token_bucket_avl_per_sec(2, 0, SHAPING_DIR_OUT, PROFILE_LIMIT_DIRECTION_INCOMING_OUTGOING);
+    stub_set_token_bucket_avl_per_sec(3, 1000, SHAPING_DIR_OUT, PROFILE_LIMIT_DIRECTION_INCOMING_OUTGOING);
     actual_tx_queue = stub_get_tx_queue();
     shaper_rules_update(&ctx->thread_ctx[0], sf, rule_id, 1);
 
@@ -801,7 +864,7 @@ TEST(single_session_async, udp_close_before_async_exec)
     ASSERT_TRUE(sf != NULL);
 
     stub_set_matched_shaping_rules(1, rule_id, priority, profile_num, profile_id);
-    stub_set_token_bucket_avl_per_sec(0, 1000, SHAPING_DIR_OUT);
+    stub_set_token_bucket_avl_per_sec(0, 1000, SHAPING_DIR_OUT, PROFILE_LIMIT_DIRECTION_INCOMING_OUTGOING);
     stub_set_async_token_get_times(0, 20);
     actual_tx_queue = stub_get_tx_queue();
     shaper_rules_update(&ctx->thread_ctx[0], sf, rule_id, 1);
@@ -868,8 +931,8 @@ TEST(two_session_diff_priority_same_profile, udp_borrow_in_order)
 
     stub_set_matched_shaping_rules(2, rule_ids, prioritys, profile_nums, profile_ids);
 
-    stub_set_token_bucket_avl_per_sec(1, 0, SHAPING_DIR_OUT);
-    stub_set_token_bucket_avl_per_sec(2, 1000, SHAPING_DIR_OUT);
+    stub_set_token_bucket_avl_per_sec(1, 0, SHAPING_DIR_OUT, PROFILE_LIMIT_DIRECTION_INCOMING_OUTGOING);
+    stub_set_token_bucket_avl_per_sec(2, 1000, SHAPING_DIR_OUT, PROFILE_LIMIT_DIRECTION_INCOMING_OUTGOING);
     actual_tx_queue = stub_get_tx_queue();
     shaper_rules_update(&ctx->thread_ctx[0], sf1, rule_id1, 1);
     shaper_rules_update(&ctx->thread_ctx[0], sf2, rule_id2, 1);
@@ -976,7 +1039,7 @@ TEST(two_session_diff_priority_same_profile, two_thread_udp_tx_in_order)
 
     stub_set_matched_shaping_rules(2, rule_ids, prioritys, profile_nums, profile_id);
 
-    stub_set_token_bucket_avl_per_sec(0, 1000, SHAPING_DIR_OUT);
+    stub_set_token_bucket_avl_per_sec(0, 1000, SHAPING_DIR_OUT, PROFILE_LIMIT_DIRECTION_INCOMING_OUTGOING);
     actual_tx_queue = stub_get_tx_queue();
     shaper_rules_update(&ctx->thread_ctx[0], sf1, rule_id1, 1);
     shaper_rules_update(&ctx->thread_ctx[1], sf2, rule_id2, 1);
@@ -1062,7 +1125,7 @@ TEST(two_session_diff_priority_same_profile, profile_timer_test)
 
     stub_set_matched_shaping_rules(2, rule_ids, prioritys, profile_nums, profile_id);
 
-    stub_set_token_bucket_avl_per_sec(0, 1000, SHAPING_DIR_OUT);
+    stub_set_token_bucket_avl_per_sec(0, 1000, SHAPING_DIR_OUT, PROFILE_LIMIT_DIRECTION_INCOMING_OUTGOING);
     actual_tx_queue = stub_get_tx_queue();
     shaper_rules_update(&ctx->thread_ctx[0], sf1, rule_id1, 1);
     shaper_rules_update(&ctx->thread_ctx[1], sf2, rule_id2, 1);
@@ -1168,8 +1231,8 @@ TEST(two_sessions, priority_non_block)
 
     stub_set_matched_shaping_rules(3, rule_ids, prioritys, profile_nums, profile_id);
 
-    stub_set_token_bucket_avl_per_sec(0, 3000, SHAPING_DIR_OUT);
-    stub_set_token_bucket_avl_per_sec(1, 1000, SHAPING_DIR_OUT);
+    stub_set_token_bucket_avl_per_sec(0, 3000, SHAPING_DIR_OUT, PROFILE_LIMIT_DIRECTION_INCOMING_OUTGOING);
+    stub_set_token_bucket_avl_per_sec(1, 1000, SHAPING_DIR_OUT, PROFILE_LIMIT_DIRECTION_INCOMING_OUTGOING);
     actual_tx_queue = stub_get_tx_queue();
     shaper_rules_update(&ctx->thread_ctx[0], sf1, rule_id1, 2);
     shaper_rules_update(&ctx->thread_ctx[1], sf2, rule_id2, 1);
@@ -1241,8 +1304,8 @@ TEST(two_sessions, borrow_when_primary_profile_priority_blocked)
 
     stub_set_matched_shaping_rules(2, rule_ids, prioritys, profile_nums, profile_id);
 
-    stub_set_token_bucket_avl_per_sec(0, 1000, SHAPING_DIR_OUT);
-    stub_set_token_bucket_avl_per_sec(1, 1000, SHAPING_DIR_OUT);
+    stub_set_token_bucket_avl_per_sec(0, 1000, SHAPING_DIR_OUT, PROFILE_LIMIT_DIRECTION_INCOMING_OUTGOING);
+    stub_set_token_bucket_avl_per_sec(1, 1000, SHAPING_DIR_OUT, PROFILE_LIMIT_DIRECTION_INCOMING_OUTGOING);
     actual_tx_queue = stub_get_tx_queue();
     shaper_rules_update(&ctx->thread_ctx[0], sf1, rule_id1, 1);
     shaper_rules_update(&ctx->thread_ctx[1], sf2, rule_id2, 1);
@@ -1316,8 +1379,8 @@ TEST(two_sessions, primary_profile_priority_blocked_by_borrow_profile)
 
     stub_set_matched_shaping_rules(2, rule_ids, prioritys, profile_nums, profile_id);
 
-    stub_set_token_bucket_avl_per_sec(0, 0, SHAPING_DIR_OUT);
-    stub_set_token_bucket_avl_per_sec(1, 1000, SHAPING_DIR_OUT);
+    stub_set_token_bucket_avl_per_sec(0, 0, SHAPING_DIR_OUT, PROFILE_LIMIT_DIRECTION_INCOMING_OUTGOING);
+    stub_set_token_bucket_avl_per_sec(1, 1000, SHAPING_DIR_OUT, PROFILE_LIMIT_DIRECTION_INCOMING_OUTGOING);
     actual_tx_queue = stub_get_tx_queue();
     shaper_rules_update(&ctx->thread_ctx[0], sf1, rule_id1, 1);
     shaper_rules_update(&ctx->thread_ctx[1], sf2, rule_id2, 1);
@@ -1376,7 +1439,7 @@ TEST(statistics, udp_drop_pkt)
     ASSERT_TRUE(sf != NULL);
     
     stub_set_matched_shaping_rules(1, rule_id, priority, profile_num, profile_id);
-    stub_set_token_bucket_avl_per_sec(0, 1000, SHAPING_DIR_OUT);
+    stub_set_token_bucket_avl_per_sec(0, 1000, SHAPING_DIR_OUT, PROFILE_LIMIT_DIRECTION_INCOMING_OUTGOING);
     actual_tx_queue = stub_get_tx_queue();
     shaper_rules_update(&ctx->thread_ctx[0], sf, rule_id, 1);
 
@@ -1452,7 +1515,7 @@ TEST(statistics, udp_queueing_pkt)
     ASSERT_TRUE(sf != NULL);
 
     stub_set_matched_shaping_rules(1, rule_id, priority, profile_num, profile_id);
-    stub_set_token_bucket_avl_per_sec(0, 1000, SHAPING_DIR_OUT);
+    stub_set_token_bucket_avl_per_sec(0, 1000, SHAPING_DIR_OUT, PROFILE_LIMIT_DIRECTION_INCOMING_OUTGOING);
     actual_tx_queue = stub_get_tx_queue();
     shaper_rules_update(&ctx->thread_ctx[0], sf, rule_id, 1);
 
diff --git a/shaping/test/stub.cpp b/shaping/test/stub.cpp
index 93ce156..8ba25f9 100644
--- a/shaping/test/stub.cpp
+++ b/shaping/test/stub.cpp
@@ -14,6 +14,7 @@
 #include <string.h>
 #include <assert.h>
 #include <pthread.h>
+#include "shaper.h"
 #include "stub.h"
 #include "shaper_maat.h"
 #include "log.h"
@@ -40,6 +41,7 @@ struct stub_token_thread_arg {
 struct stub_avaliable_token {
     int in_limit_bandwidth;
     int out_limit_bandwidth;
+    int bidirection_limit_bandwidth;
 };
 
 struct stub_pkt_queue tx_queue;
@@ -65,14 +67,21 @@ void * stub_get_token_thread_func(void *data)
     return NULL;
 }
 
-void stub_set_token_bucket_avl_per_sec(int profile_id, unsigned int tokens, unsigned char direction)
+void stub_set_token_bucket_avl_per_sec(int profile_id, unsigned int tokens, unsigned char direction, enum shaping_profile_limit_direction limit_direction)
 {
-    if (direction == SHAPING_DIR_IN) {
-        pf_array[profile_id].in_limit_bandwidth = tokens * 8;
-        pf_curr_avl_token[profile_id].in_limit_bandwidth = tokens * 8;
+    pf_array[profile_id].limit_direction = limit_direction;
+
+    if (limit_direction == PROFILE_LIMIT_DIRECTION_BIDIRECTION) {
+        pf_array[profile_id].bidirection_limit_bandwidth = tokens * 8;
+        pf_curr_avl_token[profile_id].bidirection_limit_bandwidth = tokens * 8;
     } else {
-        pf_array[profile_id].out_limit_bandwidth = tokens * 8;
-        pf_curr_avl_token[profile_id].out_limit_bandwidth = tokens * 8;
+        if (direction == SHAPING_DIR_IN) {
+            pf_array[profile_id].in_limit_bandwidth = tokens * 8;
+            pf_curr_avl_token[profile_id].in_limit_bandwidth = tokens * 8;
+        } else {
+            pf_array[profile_id].out_limit_bandwidth = tokens * 8;
+            pf_curr_avl_token[profile_id].out_limit_bandwidth = tokens * 8;
+        }
     }
 
     return;
@@ -80,6 +89,7 @@ void stub_set_token_bucket_avl_per_sec(int profile_id, unsigned int tokens, unsi
 
 void stub_refresh_token_bucket(int profile_id)
 {
+    pf_curr_avl_token[profile_id].bidirection_limit_bandwidth = pf_array[profile_id].bidirection_limit_bandwidth;
     pf_curr_avl_token[profile_id].in_limit_bandwidth = pf_array[profile_id].in_limit_bandwidth;
     pf_curr_avl_token[profile_id].out_limit_bandwidth = pf_array[profile_id].out_limit_bandwidth;
     return;
@@ -191,6 +201,7 @@ void stub_init()
     for (i = 0; i < MAX_STUB_PROFILE_NUM; i++) {
         pf_curr_avl_token[i].in_limit_bandwidth = DEFAULT_AVALIABLE_TOKEN_PER_SEC;
         pf_curr_avl_token[i].out_limit_bandwidth = DEFAULT_AVALIABLE_TOKEN_PER_SEC;
+        pf_curr_avl_token[i].bidirection_limit_bandwidth = DEFAULT_AVALIABLE_TOKEN_PER_SEC;
         pf_array[i].id = i;
         pf_array[i].in_limit_bandwidth = DEFAULT_AVALIABLE_TOKEN_PER_SEC;
         pf_array[i].out_limit_bandwidth = DEFAULT_AVALIABLE_TOKEN_PER_SEC;
@@ -377,7 +388,14 @@ void swarmkv_tconsume(struct swarmkv * db, const char * key, size_t keylen, long
 
     sscanf(key, "tsg-shaping-%d-%15s", &profile_id, direction);
 
-    if (strncmp("incoming", direction, sizeof(direction)) == 0) {
+    if (strncmp("bidirectional", direction, sizeof(direction)) == 0) {
+        if (pf_curr_avl_token[profile_id].bidirection_limit_bandwidth == DEFAULT_AVALIABLE_TOKEN_PER_SEC) {
+            actual_tokens = tokens;
+        } else {
+            actual_tokens = pf_curr_avl_token[profile_id].bidirection_limit_bandwidth >= tokens ? tokens : 0;
+            pf_curr_avl_token[profile_id].bidirection_limit_bandwidth -= actual_tokens;
+        }
+    } else if (strncmp("incoming", direction, sizeof(direction)) == 0) {
         if (pf_curr_avl_token[profile_id].in_limit_bandwidth == DEFAULT_AVALIABLE_TOKEN_PER_SEC) {
             actual_tokens = tokens;
         } else {
diff --git a/shaping/test/stub.h b/shaping/test/stub.h
index e4c0476..a00bd01 100644
--- a/shaping/test/stub.h
+++ b/shaping/test/stub.h
@@ -30,7 +30,7 @@ struct stub_packet_node {
 TAILQ_HEAD(stub_pkt_queue, stub_packet_node);
 
 
-void stub_set_token_bucket_avl_per_sec(int profile_id, unsigned int tokens, unsigned char direction);
+void stub_set_token_bucket_avl_per_sec(int profile_id, unsigned int tokens, unsigned char direction, enum shaping_profile_limit_direction limit_direction);
 void stub_refresh_token_bucket(int profile_id);
 void stub_set_profile_type(int profile_id, enum shaping_profile_type type);
 void stub_set_async_token_get_times(int profile_id, int times);