Add responseless dns requests to dns_data table. pixie-io#613

Summary:
This pull request is related to pixie-io#418 which relates to including those DNS requests into the dns_data table for which a matching response was not found.

Closes pixie-io#613.

Test Plan: Verified that tests in stitcher_test.cc are passed.

Reviewers: #stirling, jps

Reviewed By: #stirling, jps

Signed-off-by: Noman_ali_bajwa <nmnbajwa@gmail.com>

Differential Revision: https://phab.corp.pixielabs.ai/D12442

GitOrigin-RevId: 934ec53aeafdd1794638d7f022248becd1d2acfe

src/stirling/source_connectors/socket_tracer/protocols/dns/BUILD.bazel

@@ -48,5 +48,8 @@ pl_cc_test(
 pl_cc_test(
     name = "stitcher_test",
     srcs = ["stitcher_test.cc"],
-    deps = [":cc_library"],
+    deps = [
+        ":cc_library",
+        "//src/stirling/testing:cc_library",
+    ],
 )

src/stirling/source_connectors/socket_tracer/protocols/dns/stitcher.cc

@@ -22,13 +22,25 @@
 #include <rapidjson/stringbuffer.h>
 #include <rapidjson/writer.h>
 
+#include <chrono>
 #include <deque>
 #include <string>
 #include <utility>
 
 #include "src/common/base/base.h"
 #include "src/stirling/source_connectors/socket_tracer/protocols/dns/types.h"
 
+DEFINE_bool(include_respless_dns_requests, false,
+            "If true, use customStitchFrames otherwise uses simple StitchFrames");
+
+DEFINE_uint64(dns_request_timeout_threshold_milliseconds, 2000,
+              "Number of seconds to wait for the in-flight response of a dns request. Depends on "
+              "include_respless_dns_requests.");
+
+auto current_time = std::chrono::duration_cast<std::chrono::milliseconds>(
+                        std::chrono::system_clock::now().time_since_epoch())
+                        .count();
+
 namespace px {
 namespace stirling {
 namespace protocols {
@@ -176,7 +188,8 @@ StatusOr<Record> ProcessReqRespPair(const Frame& req_frame, const Frame& resp_fr
 // For each response that is at the head of the deque, there should exist a previous request with
 // the same txid. Find it, and consume both frames.
 RecordsWithErrorCount<Record> StitchFrames(std::deque<Frame>* req_frames,
-                                           std::deque<Frame>* resp_frames) {
+                                           std::deque<Frame>* resp_frames,
+                                           bool include_respless_dns_requests) {
   std::vector<Record> entries;
   int error_count = 0;
 
@@ -237,7 +250,31 @@ RecordsWithErrorCount<Record> StitchFrames(std::deque<Frame>* req_frames,
 
   resp_frames->clear();
 
-  return {entries, error_count};
+  if (include_respless_dns_requests) {
+    // After the external loop's lifecycle comes to an end we end up with the request deque
+    // having only those request frames which have not been consumed yet i.e. consumed = false
+    // so essentially these are the requests which could not be matched with any response frame.
+    // Hence we iterate over this request deque, add a default response to it, make a record and
+    // append those records at the end of the entries vector.
+    auto it = req_frames->begin();
+    while (it != req_frames->end()) {
+      if (!(it->consumed)) {
+        auto elapsed_seconds = current_time - it->timestamp_ns;
+
+        if (elapsed_seconds > FLAGS_dns_request_timeout_threshold_milliseconds) {
+          Frame default_resp_frame;
+          default_resp_frame.timestamp_ns =
+              (it->timestamp_ns) + FLAGS_dns_request_timeout_threshold_milliseconds;
+          StatusOr<Record> record_status = ProcessReqRespPair(*it, default_resp_frame);
+          entries.push_back(record_status.ConsumeValueOrDie());
+        }
+      }
+      it++;
+    }
+    return {entries, error_count};
+  } else {
+    return {entries, error_count};
+  }
 }
 
 }  // namespace dns

src/stirling/source_connectors/socket_tracer/protocols/dns/stitcher.h

@@ -26,6 +26,11 @@
 #include "src/stirling/source_connectors/socket_tracer/protocols/common/interface.h"
 #include "src/stirling/source_connectors/socket_tracer/protocols/dns/types.h"
 
+// Choose either the simple implementation of StitchFrames (without missing records) or
+// to use CustomStitchFrames for including responseless DNS requests as well.
+DECLARE_bool(include_respless_dns_requests);
+DECLARE_uint64(dns_request_timeout_threshold_milliseconds);
+
 namespace px {
 namespace stirling {
 namespace protocols {
@@ -36,18 +41,21 @@ namespace dns {
  *
  * @param req_packets: deque of all request frames.
  * @param resp_packets: deque of all response frames.
+ * @param include_respless_dns_requests: bool to decide whether to include responseless DNS requests
+ * or not.
  * @return A vector of entries to be appended to table store.
  */
 RecordsWithErrorCount<Record> StitchFrames(std::deque<Frame>* req_packets,
-                                           std::deque<Frame>* resp_packets);
+                                           std::deque<Frame>* resp_packets,
+                                           bool include_respless_dns_requests);
 
 }  // namespace dns
 
 template <>
 inline RecordsWithErrorCount<dns::Record> StitchFrames(std::deque<dns::Frame>* req_packets,
                                                        std::deque<dns::Frame>* resp_packets,
                                                        NoState* /* state */) {
-  return dns::StitchFrames(req_packets, resp_packets);
+  return dns::StitchFrames(req_packets, resp_packets, FLAGS_include_respless_dns_requests);
 }
 
 }  // namespace protocols

src/stirling/source_connectors/socket_tracer/protocols/dns/stitcher_test.cc

@@ -23,6 +23,7 @@
 
 #include "src/common/base/inet_utils.h"
 #include "src/stirling/source_connectors/socket_tracer/protocols/dns/stitcher.h"
+#include "src/stirling/testing/common.h"
 
 using ::testing::HasSubstr;
 using ::testing::IsEmpty;
@@ -90,7 +91,7 @@ TEST(DnsStitcherTest, RecordOutput) {
   req_frames.push_back(req0_frame);
   resp_frames.push_back(resp0_frame);
 
-  result = StitchFrames(&req_frames, &resp_frames);
+  result = StitchFrames(&req_frames, &resp_frames, FLAGS_include_respless_dns_requests);
   EXPECT_TRUE(resp_frames.empty());
   EXPECT_EQ(req_frames.size(), 0);
   EXPECT_EQ(result.error_count, 0);
@@ -116,6 +117,8 @@ TEST(DnsStitcherTest, OutOfOrderMatching) {
   std::deque<Frame> resp_frames;
   RecordsWithErrorCount<Record> result;
 
+  PL_SET_FOR_SCOPE(FLAGS_include_respless_dns_requests, true);
+
   int t = 0;
 
   Frame req0_frame = CreateReqFrame(++t, 0);
@@ -125,7 +128,7 @@ TEST(DnsStitcherTest, OutOfOrderMatching) {
   Frame req2_frame = CreateReqFrame(++t, 2);
   Frame resp2_frame = CreateRespFrame(++t, 2, std::vector<DNSRecord>());
 
-  result = StitchFrames(&req_frames, &resp_frames);
+  result = StitchFrames(&req_frames, &resp_frames, FLAGS_include_respless_dns_requests);
   EXPECT_TRUE(resp_frames.empty());
   EXPECT_EQ(req_frames.size(), 0);
   EXPECT_EQ(result.error_count, 0);
@@ -134,38 +137,38 @@ TEST(DnsStitcherTest, OutOfOrderMatching) {
   req_frames.push_back(req0_frame);
   req_frames.push_back(req1_frame);
 
-  result = StitchFrames(&req_frames, &resp_frames);
+  result = StitchFrames(&req_frames, &resp_frames, FLAGS_include_respless_dns_requests);
   EXPECT_TRUE(resp_frames.empty());
   EXPECT_EQ(req_frames.size(), 2);
   EXPECT_EQ(result.error_count, 0);
-  EXPECT_EQ(result.records.size(), 0);
+  EXPECT_EQ(result.records.size(), FLAGS_include_respless_dns_requests ? 2 : 0);
 
   resp_frames.push_back(resp1_frame);
 
-  result = StitchFrames(&req_frames, &resp_frames);
+  result = StitchFrames(&req_frames, &resp_frames, FLAGS_include_respless_dns_requests);
   EXPECT_TRUE(resp_frames.empty());
   EXPECT_EQ(req_frames.size(), 2);
   EXPECT_EQ(result.error_count, 0);
-  EXPECT_EQ(result.records.size(), 1);
+  EXPECT_EQ(result.records.size(), FLAGS_include_respless_dns_requests ? 2 : 1);
 
   req_frames.push_back(req2_frame);
   resp_frames.push_back(resp0_frame);
 
-  result = StitchFrames(&req_frames, &resp_frames);
+  result = StitchFrames(&req_frames, &resp_frames, FLAGS_include_respless_dns_requests);
   EXPECT_TRUE(resp_frames.empty());
   EXPECT_EQ(req_frames.size(), 1);
   EXPECT_EQ(result.error_count, 0);
-  EXPECT_EQ(result.records.size(), 1);
+  EXPECT_EQ(result.records.size(), FLAGS_include_respless_dns_requests ? 2 : 1);
 
   resp_frames.push_back(resp2_frame);
 
-  result = StitchFrames(&req_frames, &resp_frames);
+  result = StitchFrames(&req_frames, &resp_frames, FLAGS_include_respless_dns_requests);
   EXPECT_TRUE(resp_frames.empty());
   EXPECT_EQ(resp_frames.size(), 0);
   EXPECT_EQ(result.error_count, 0);
   EXPECT_EQ(result.records.size(), 1);
 
-  result = StitchFrames(&req_frames, &resp_frames);
+  result = StitchFrames(&req_frames, &resp_frames, FLAGS_include_respless_dns_requests);
   EXPECT_TRUE(resp_frames.empty());
   EXPECT_EQ(resp_frames.size(), 0);
   EXPECT_EQ(result.error_count, 0);