safe add

3 files changed, 33 insertions, 18 deletions

diff --git a/src/metrics/metric.c b/src/metrics/metric.c
index 887102a..95b7200 100644
--- a/src/metrics/metric.c
+++ b/src/metrics/metric.c
@@ -4,6 +4,8 @@
 #include <assert.h>
 #include <string.h>
 #include <stdarg.h>
+#include <limits.h>
+
 
 #include "uthash.h"
 #include "histogram_encoder.h"
@@ -711,11 +713,22 @@ struct metric *metric_counter_new(const char *name)
     return pthis;
 }
 
+long long safe_add_longlong(long long a, long long b)
+{
+    if (a > 0 && b > LLONG_MAX - a) {
+        return LLONG_MAX;
+    }
+    if (a < 0 && b < LLONG_MIN - a) {
+        return LLONG_MIN;
+    }
+    return a + b;
+}
+
 void metric_counter_incrby(struct metric *pthis, int cell_id, long long value)
 {
     struct metric_measure_data *data = metric_find_or_new_cell(pthis, cell_id);
     struct metric_counter_or_gauge *counter = data->counter;
-    counter->value += value;
+    counter->value = safe_add_longlong(counter->value, value);
 }
 
 int metric_counter_set(struct metric *pthis, int cell_id, long long value)
diff --git a/src/tags/heavy_keeper.c b/src/tags/heavy_keeper.c
index f866484..a82677b 100644
--- a/src/tags/heavy_keeper.c
+++ b/src/tags/heavy_keeper.c
@@ -9,6 +9,7 @@
 #include <stdlib.h>
 #include <string.h>
 #include <assert.h>
+#include <limits.h>
 
 #include "minheap/heap.h"
 #include "mpack/mpack.h"
@@ -157,8 +158,8 @@ struct Bucket *map_flow_id_hash_to_bucket(struct heavy_keeper *hk, int array_ind
 }
 
 unsigned long long count_add(unsigned long long a, unsigned long long b) {
-    if (a > UINT64_MAX - b) {
-        return UINT64_MAX;
+    if (a > LLONG_MAX - b) {
+        return LLONG_MAX;
     }
     return a + b;
 }
@@ -335,7 +336,7 @@ static void heavy_keeper_merge_sketch(struct heavy_keeper *dest, const struct he
             const struct Bucket *bucket_src = &(src->sketch[array_id * w + bucket_id]);
 
             if (bucket_dest->finger_print == bucket_src->finger_print) {
-                bucket_dest->count += bucket_src->count;
+                bucket_dest->count = count_add(bucket_dest->count, bucket_src->count);
             } else {
                 if (bucket_dest->count < bucket_src->count) {  // bucket_src stores the elephant flow.
                     bucket_dest->count = bucket_src->count;
@@ -433,7 +434,7 @@ void heavy_keeper_merge_sorted_set_recording_id_details(struct heavy_keeper *des
     int cell_id_in_src_before_merge_len = 0;
     int cell_id_in_dest_corresponding[dest->K]; // length == cell_id_in_src_before_merge_len
 
-    /* ------------------------------ merge dest ------------------------------ */
+    /* ------------------------------ merge dest ------------------------------ *
     for (int i = 0; i < size_dest; i++) {
         unsigned long long maxv = find_maxv_in_sketch(dest, tag_arr[i]);
         assert(tag_arr[i] != NULL);
diff --git a/src/tags/sorted_set.c b/src/tags/sorted_set.c
index df6002e..6d71777 100644
--- a/src/tags/sorted_set.c
+++ b/src/tags/sorted_set.c
@@ -3,6 +3,7 @@
 #include <string.h>

 #include <stdio.h>

 #include <assert.h>

+#include <limits.h>

 

 #include "my_ut_hash_inner.h"

 #include "xxhash/xxhash.h"

@@ -210,14 +211,17 @@ bool sorted_set_check_is_full(const struct sorted_set *ss)
     return ss->heap->cur_size >= ss->heap->max_size;

 }

 

-void sorted_set_insert_to_available_heap(struct sorted_set *ss, const struct tag_hash_key *tag, unsigned long long cnt, int cell_id)

+unsigned long long safe_add(unsigned long long a, unsigned long long b)

 {

-    // cnt += 1; // sorted set will let the count start from 1, 0 for dying entry.

-    if (cnt >= UINT64_MAX) {

-        cnt = UINT64_MAX;

-    } else {

-        cnt += 1; // sorted set will let the count start from 1, 0 for dying entry.

+    if (ULLONG_MAX - a < b) {

+        return ULLONG_MAX;

     }

+    return a + b;

+}

+

+void sorted_set_insert_to_available_heap(struct sorted_set *ss, const struct tag_hash_key *tag, unsigned long long cnt, int cell_id)

+{

+    cnt = safe_add(cnt, 1); // sorted set will let the count start from 1, 0 for dying entry.

     unsigned long long *tmp_cnt = (unsigned long long*)malloc(sizeof(unsigned long long));

     *tmp_cnt = cnt;

     struct entry_data *tmp_data = entry_data_construct(tag, cell_id);

@@ -267,12 +271,8 @@ int sorted_set_insert_and_record_popped_back(struct sorted_set *ss, const struct
         if (!sorted_set_entry_dying(entry)) {

             return -1;

         }

-        if (cnt >= UINT64_MAX) {

-            cnt = UINT64_MAX;

-        } else {

-            cnt += 1; // sorted set will let the count start from 1, 0 for dying entry.

-        }

-        sorted_set_entry_set_key(ss, entry, cnt); // count 0 for dying entry.

+        

+        sorted_set_entry_set_key(ss, entry, safe_add(cnt, 1)); // sorted set will let the count start from 1, 0 for dying entry.

         struct entry_data *data = sorted_set_entry_get_data(entry);

         data->cell_id = cell_id;

         ss->n_living_entry++;

@@ -350,7 +350,8 @@ int sorted_set_incrby(struct sorted_set *ss, const struct tag_hash_key *tag, uns
     if (entry == NULL) {

         return -1;

     }

-    sorted_set_entry_set_key(ss, entry, sorted_set_entry_get_count(entry) + count);

+    unsigned long long cnt_old = sorted_set_entry_get_count(entry);

+    sorted_set_entry_set_key(ss, entry, safe_add(count, cnt_old));

     return 0;

 }