Merge pull request #658 from bytedance/fix-python

feat python&node discard count
bytedance · Jul 17, 2024 · 63aa263 · 63aa263
2 parents b38e48b + 7217493
commit 63aa263
Show file tree

Hide file tree

Showing 8 changed files with 49 additions and 15 deletions.
diff --git a/rasp/node/src/client.js b/rasp/node/src/client.js
@@ -126,12 +126,16 @@ class SmithClient extends EventEmitter {
         });
     }
 
-    postMessage(operate, data) {
-        if (!this._connected)
+    postMessage(operate, data, discard_send = () => {}) {
+        if (!this._connected) {
+            discard_send();
             return;
+        }
 
-        if ('writableLength' in this._socket && this._socket.writableLength > BUFFER_MAX_SIZE)
+        if ('writableLength' in this._socket && this._socket.writableLength > BUFFER_MAX_SIZE) {
+            discard_send();
             return;
+        }
 
         const message = {
             'pid': process.pid,
@@ -151,7 +155,9 @@ class SmithClient extends EventEmitter {
         buffer.writeUInt32BE(length, 0);
         buffer.write(payload, 4);
 
-        this._socket.write(buffer);
+        if (!this._socket.write(buffer)) {
+            discard_send();
+        }
     }
 }
 

diff --git a/rasp/node/src/smith.js b/rasp/node/src/smith.js
@@ -13,6 +13,7 @@ const LOGICAL_OR = 0;
 const LOGICAL_AND = 1;
 
 const heartbeat = {};
+heartbeat.discard_send = 0;
 
 const blocks = new Map();
 const filters = new Map();
@@ -168,15 +169,19 @@ function smithHook(func, classID, methodID, canBlock = false, processors = {}) {
 
                 return false;
             })) {
-                client.postMessage(Operate.trace, smithTrace);
+                client.postMessage(Operate.trace, smithTrace, function() {
+                    heartbeat.discard_send++;
+                });
                 throw new Error('API blocked by RASP');
             }
         }
 
         const filter = filters.get(`${classID} ${methodID}`);
 
         if (!filter) {
-            client.postMessage(Operate.trace, smithTrace);
+            client.postMessage(Operate.trace, smithTrace, function() {
+                heartbeat.discard_send++;
+            });
             return func.call(this, ...args);
         }
 
@@ -191,7 +196,9 @@ function smithHook(func, classID, methodID, canBlock = false, processors = {}) {
             return func.call(this, ...args);
         }
 
-        client.postMessage(Operate.trace, smithTrace);
+        client.postMessage(Operate.trace, smithTrace, function() {
+            heartbeat.discard_send++;
+        });
 
         return func.call(this, ...args);
     }

diff --git a/rasp/python/python-probe/rasp/probe/client/smith_client.cpp b/rasp/python/python-probe/rasp/probe/client/smith_client.cpp
@@ -1,4 +1,5 @@
 #include "smith_client.h"
+#include "smith_probe.h"
 #include <aio/ev/timer.h>
 #include <aio/net/stream.h>
 #include <zero/log.h>
@@ -80,7 +81,7 @@ startClient(const std::shared_ptr<aio::Context> &context) {
                     reason.code,
                     reason.message.c_str()
             );
-
+            gProbe->discard_send++;
             return zero::ptr::makeRef<aio::ev::Timer>(context)->setTimeout(1min);
         });
     });

diff --git a/rasp/python/python-probe/rasp/probe/client/smith_message.cpp b/rasp/python/python-probe/rasp/probe/client/smith_message.cpp
@@ -41,7 +41,10 @@ void to_json(nlohmann::json &j, const Heartbeat &heartbeat) {
     j = {
             {"filter", heartbeat.filter},
             {"block",  heartbeat.block},
-            {"limit",  heartbeat.limit}
+            {"limit",  heartbeat.limit},
+            {"discard_surplus", heartbeat.discard_surplus},
+            {"discard_send", heartbeat.discard_send},
+            {"discard_post", heartbeat.discard_post}
     };
 }
 

diff --git a/rasp/python/python-probe/rasp/probe/client/smith_message.h b/rasp/python/python-probe/rasp/probe/client/smith_message.h
@@ -38,6 +38,9 @@ struct Heartbeat {
     std::string filter;
     std::string block;
     std::string limit;
+    int64_t discard_surplus;
+    int64_t discard_send;
+    int64_t discard_post;
 };
 
 struct Trace {

diff --git a/rasp/python/python-probe/rasp/probe/client/smith_probe.cpp b/rasp/python/python-probe/rasp/probe/client/smith_probe.cpp
@@ -169,7 +169,9 @@ void startProbe() {
                 gProbe->quotas[cid][mid] = it->second;
             }
         }
-
+        heartbeat->discard_post = gProbe->discard_post;
+        heartbeat->discard_send = gProbe->discard_send;
+        heartbeat->discard_surplus = gProbe->discard_surplus;
         sender->trySend({HEARTBEAT, *heartbeat});
         return true;
     });
@@ -370,8 +372,12 @@ void startProbe() {
                 Trace trace = gProbe->buffer[*index];
                 gProbe->buffer.release(*index);
 
-                if (pass(trace, *filters))
-                    sender->trySend({TRACE, gProbe->buffer[*index]});
+                if (pass(trace, *filters)) {
+                    auto result = sender->trySend({TRACE, gProbe->buffer[*index]});
+                    if (!result) {
+                        gProbe->discard_send++;
+                    }
+                }
 
                 P_CONTINUE(loop);
             }
@@ -402,7 +408,8 @@ void startProbe() {
                 ExceptionInfo info = gProbe->info[*index];
                 gProbe->info.release(*index);
 
-                sender->trySend({EXCEPTIONINFO, gProbe->info[*index]});
+                if(!sender->trySend({EXCEPTIONINFO, gProbe->info[*index]}))
+                    gProbe->discard_send++;
 
                 P_CONTINUE(loop);
             }

diff --git a/rasp/python/python-probe/rasp/probe/client/smith_probe.h b/rasp/python/python-probe/rasp/probe/client/smith_probe.h
@@ -41,6 +41,9 @@ struct Probe {
     zero::atomic::CircularBuffer<Trace, TRACE_BUFFER_SIZE> buffer;
     zero::atomic::CircularBuffer<ExceptionInfo, EXCEPTIONINFO_BUFFER_SIZE> info;
     zero::atomic::CircularBuffer<Policy *, PREPARED_POLICY_COUNT> nodes;
+    std::atomic<int64_t> discard_surplus;
+    std::atomic<int64_t> discard_send;
+    std::atomic<int64_t> discard_post;
 
     Policy *popNode();
     bool pushNode(Policy *node);

diff --git a/rasp/python/python-probe/rasp/probe/library.cpp b/rasp/python/python-probe/rasp/probe/library.cpp
@@ -70,8 +70,10 @@ PyObject *send(PyObject *self, PyObject *args) {
 
     std::optional<size_t> index = gProbe->buffer.reserve();
 
-    if (!index)
+    if (!index) {
+        gProbe->discard_post++;
         Py_RETURN_NONE;
+    }
 
     gProbe->buffer[*index] = *(Trace *) pyTrace;
     gProbe->buffer.commit(*index);
@@ -167,8 +169,10 @@ PyObject *surplus(PyObject *self, PyObject *args) {
     int n = quota;
 
     do {
-        if (n <= 0)
+        if (n <= 0) {
+            gProbe->discard_surplus++;
             Py_RETURN_FALSE;
+        }
     } while (!quota.compare_exchange_weak(n, n - 1));
 
     Py_RETURN_TRUE;