Skip to content

Commit 0e6ef3e

Browse files
kpattaswamykpattaswamy
authored
Add the BPF test for Mongo protocol tracing (#1778)
Summary: This PR adds the BPF test for the mongo protocol tracer. The default tracing mode is on. Related issues: #640 Type of change: /kind feature Test Plan: Added the BPF test Changelog Message: ``` MongoDB query profiling is now supported by Stirling. ``` Signed-off-by: Kartik Pattaswamy <kpattaswamy@pixielabs.ai>
1 parent b1aa1a6 commit 0e6ef3e

File tree

17 files changed

+484
-3
lines changed

17 files changed

+484
-3
lines changed

‎WORKSPACE‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -297,3 +297,12 @@ maven_install(
297297
load("@px_deps//:defs.bzl", px_deps_pinned_maven_install = "pinned_maven_install")
298298

299299
px_deps_pinned_maven_install()
300+
301+
pip_parse(
302+
name = "mongodb_bpf_test_requirements",
303+
requirements_lock = "//src/stirling/source_connectors/socket_tracer/testing/containers/mongodb:requirements.bazel.txt",
304+
)
305+
306+
load("@mongodb_bpf_test_requirements//:requirements.bzl", mongodb_bpf_test_install_deps = "install_deps")
307+
308+
mongodb_bpf_test_install_deps()

‎bazel/container_images.bzl‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -278,3 +278,11 @@ def stirling_test_images():
278278
repository = "google-samples/microservices-demo/emailservice",
279279
digest = "sha256:d42ee712cbb4806a8b922e303a5e6734f342dfb6c92c81284a289912165b7314",
280280
)
281+
282+
# Tag: mongo:7.0
283+
# Arch: linux/amd64
284+
_container_image(
285+
name = "mongo_7_0",
286+
repository = "mongo",
287+
digest = "sha256:19b2e5c91f92c7b18113a1501c5a5fe52b71a6c6d2a5232eeebb4f2abacae04a",
288+
)

‎src/stirling/source_connectors/socket_tracer/BUILD.bazel‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -579,3 +579,18 @@ pl_cc_bpf_test(
579579
"//src/stirling/source_connectors/socket_tracer/testing/container_images:rabbitmq_producer_container",
580580
],
581581
)
582+
583+
pl_cc_bpf_test(
584+
name = "mongodb_trace_bpf_test",
585+
timeout = "moderate",
586+
srcs = ["mongodb_trace_bpf_test.cc"],
587+
tags = ["requires_bpf"],
588+
deps = [
589+
":cc_library",
590+
"//src/common/testing/test_utils:cc_library",
591+
"//src/stirling/source_connectors/socket_tracer/testing:cc_library",
592+
"//src/stirling/source_connectors/socket_tracer/testing/container_images:mongodb_client_container",
593+
"//src/stirling/source_connectors/socket_tracer/testing/container_images:mongodb_container",
594+
"//src/stirling/testing:cc_library",
595+
],
596+
)

‎src/stirling/source_connectors/socket_tracer/bcc_bpf/socket_trace.c‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -312,7 +312,7 @@ static __inline void update_traffic_class(struct conn_info_t* conn_info,
312312
struct protocol_message_t inferred_protocol = infer_protocol(buf, count, conn_info);
313313

314314
// Could not infer the traffic.
315-
if (inferred_protocol.protocol == kProtocolUnknown || conn_info->protocol == kProtocolMongo) {
315+
if (inferred_protocol.protocol == kProtocolUnknown) {
316316
return;
317317
}
318318

‎src/stirling/source_connectors/socket_tracer/mongodb_trace_bpf_test.cc‎

Lines changed: 140 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,140 @@
1+
/*
2+
* Copyright 2018- The Pixie Authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*
16+
* SPDX-License-Identifier: Apache-2.0
17+
*/
18+
19+
#include <gmock/gmock.h>
20+
#include <gtest/gtest.h>
21+
22+
#include <string>
23+
24+
#include <absl/strings/str_replace.h>
25+
26+
#include "src/common/base/base.h"
27+
#include "src/common/exec/exec.h"
28+
#include "src/common/testing/testing.h"
29+
#include "src/shared/types/column_wrapper.h"
30+
#include "src/shared/types/types.h"
31+
#include "src/stirling/core/data_table.h"
32+
#include "src/stirling/core/output.h"
33+
#include "src/stirling/source_connectors/socket_tracer/mongodb_table.h"
34+
#include "src/stirling/source_connectors/socket_tracer/protocols/mongodb/types.h"
35+
#include "src/stirling/source_connectors/socket_tracer/testing/container_images/mongodb_client_container.h"
36+
#include "src/stirling/source_connectors/socket_tracer/testing/container_images/mongodb_container.h"
37+
#include "src/stirling/source_connectors/socket_tracer/testing/protocol_checkers.h"
38+
#include "src/stirling/source_connectors/socket_tracer/testing/socket_trace_bpf_test_fixture.h"
39+
#include "src/stirling/testing/common.h"
40+
#include "src/stirling/utils/linux_headers.h"
41+
42+
namespace px {
43+
namespace stirling {
44+
45+
namespace mongodb = protocols::mongodb;
46+
47+
using ::px::stirling::testing::FindRecordIdxMatchesPID;
48+
using ::px::stirling::testing::FindRecordsMatchingPID;
49+
using ::px::stirling::testing::GetTargetRecords;
50+
using ::px::stirling::testing::SocketTraceBPFTestFixture;
51+
52+
using ::testing::AllOf;
53+
using ::testing::Eq;
54+
using ::testing::Field;
55+
using ::testing::HasSubstr;
56+
57+
void Init() {
58+
// Enable mongodb tracing.
59+
FLAGS_stirling_enable_mongodb_tracing = true;
60+
61+
// Turn off CQL and NATS tracing to give some BPF instructions back for MongoDB.
62+
// This is required for older kernels with only 4096 BPF instructions.
63+
FLAGS_stirling_enable_cass_tracing = false;
64+
FLAGS_stirling_enable_nats_tracing = false;
65+
}
66+
67+
class MongoDBTraceTest : public SocketTraceBPFTestFixture</* TClientSideTracing */ true> {
68+
protected:
69+
MongoDBTraceTest() {
70+
Init();
71+
PX_CHECK_OK(mongodb_server_.Run(std::chrono::seconds{120}));
72+
}
73+
74+
void RunMongoDBClient() {
75+
mongodb_client_.Run(
76+
std::chrono::seconds{120},
77+
{absl::Substitute("--network=container:$0", mongodb_server_.container_name())});
78+
}
79+
80+
::px::stirling::testing::MongoDBClientContainer mongodb_client_;
81+
::px::stirling::testing::MongoDBContainer mongodb_server_;
82+
};
83+
84+
auto EqMongoDBMsgType(const protocols::mongodb::Frame& f) {
85+
return Field(&protocols::mongodb::Frame::op_msg_type, Eq(f.op_msg_type));
86+
}
87+
88+
auto ContainsMongoDBMsgBody(const protocols::mongodb::Frame& f) {
89+
return Field(&protocols::mongodb::Frame::frame_body, HasSubstr(f.frame_body));
90+
}
91+
92+
auto EqMongoDBRecord(const protocols::mongodb::Record& r) {
93+
return AllOf(Field(&protocols::mongodb::Record::req, EqMongoDBMsgType(r.req)),
94+
Field(&protocols::mongodb::Record::resp, EqMongoDBMsgType(r.resp)),
95+
Field(&protocols::mongodb::Record::req, ContainsMongoDBMsgBody(r.req)),
96+
Field(&protocols::mongodb::Record::resp, ContainsMongoDBMsgBody(r.resp)));
97+
}
98+
99+
mongodb::Record RecordOpMsg(std::string req_cmd, std::string resp_status, std::string req_body,
100+
std::string resp_body) {
101+
mongodb::Record r = {};
102+
r.req.op_msg_type = req_cmd;
103+
r.req.frame_body = req_body;
104+
r.resp.op_msg_type = resp_status;
105+
r.resp.frame_body = resp_body;
106+
return r;
107+
}
108+
109+
//-----------------------------------------------------------------------------
110+
// Test Scenarios
111+
//-----------------------------------------------------------------------------
112+
113+
TEST_F(MongoDBTraceTest, Capture) {
114+
// Initiate the mongo transactions.
115+
StartTransferDataThread();
116+
RunMongoDBClient();
117+
StopTransferDataThread();
118+
119+
// Grab the data from Stirling.
120+
std::vector<TaggedRecordBatch> tablets = ConsumeRecords(SocketTraceConnector::kMongoDBTableNum);
121+
ASSERT_NOT_EMPTY_AND_GET_RECORDS(const types::ColumnWrapperRecordBatch& record_batch, tablets);
122+
123+
std::vector<mongodb::Record> server_records =
124+
GetTargetRecords<mongodb::Record>(record_batch, mongodb_server_.process_pid());
125+
126+
mongodb::Record opMsgInsert = RecordOpMsg("insert", "ok: {$numberDouble: 1.0}", "foo", "ok");
127+
mongodb::Record opMsgFind1 = RecordOpMsg("find", "cursor", "find", "foo");
128+
mongodb::Record opMsgUpdate = RecordOpMsg("update", "ok: {$numberDouble: 1.0}", "bar", "ok");
129+
mongodb::Record opMsgFind2 = RecordOpMsg("find", "cursor", "find", "bar");
130+
mongodb::Record opMsgDelete = RecordOpMsg("delete", "ok: {$numberDouble: 1.0}", "bar", "ok");
131+
132+
EXPECT_THAT(server_records, Contains(EqMongoDBRecord(opMsgInsert)));
133+
EXPECT_THAT(server_records, Contains(EqMongoDBRecord(opMsgFind1)));
134+
EXPECT_THAT(server_records, Contains(EqMongoDBRecord(opMsgUpdate)));
135+
EXPECT_THAT(server_records, Contains(EqMongoDBRecord(opMsgFind2)));
136+
EXPECT_THAT(server_records, Contains(EqMongoDBRecord(opMsgDelete)));
137+
}
138+
139+
} // namespace stirling
140+
} // namespace px

‎src/stirling/source_connectors/socket_tracer/protocols/mongodb/types.h‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -191,6 +191,9 @@ struct ProtocolTraits : public BaseProtocolTraits<Record> {
191191
};
192192

193193
} // namespace mongodb
194+
195+
template <>
196+
mongodb::stream_id_t GetStreamID(mongodb::Frame* frame);
194197
} // namespace protocols
195198
} // namespace stirling
196199
} // namespace px

‎src/stirling/source_connectors/socket_tracer/socket_trace_connector.cc‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -114,7 +114,7 @@ DEFINE_int32(stirling_enable_amqp_tracing,
114114
"If true, stirling will trace and process AMQP messages.");
115115
DEFINE_int32(stirling_enable_mongodb_tracing,
116116
gflags::Int32FromEnv("PX_STIRLING_ENABLE_MONGODB_TRACING",
117-
px::stirling::TraceMode::On),
117+
px::stirling::TraceMode::OnForNewerKernel),
118118
"If true, stirling will trace and process MongoDB messages");
119119
DEFINE_bool(stirling_disable_golang_tls_tracing,
120120
gflags::BoolFromEnv("PX_STIRLING_DISABLE_GOLANG_TLS_TRACING", false),
@@ -265,7 +265,7 @@ void SocketTraceConnector::InitProtocolTransferSpecs() {
265265
kMuxTableNum,
266266
{kRoleClient, kRoleServer},
267267
TRANSFER_STREAM_PROTOCOL(mux)}},
268-
{kProtocolMongo, TransferSpec{px::stirling::TraceMode::Off,
268+
{kProtocolMongo, TransferSpec{FLAGS_stirling_enable_mongodb_tracing,
269269
kMongoDBTableNum,
270270
{kRoleClient, kRoleServer},
271271
TRANSFER_STREAM_PROTOCOL(mongodb)}},

‎src/stirling/source_connectors/socket_tracer/testing/container_images/BUILD.bazel‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -141,6 +141,26 @@ pl_cc_test_library(
141141
deps = ["//src/common/testing/test_utils:cc_library"],
142142
)
143143

144+
pl_cc_test_library(
145+
name = "mongodb_client_container",
146+
srcs = [],
147+
hdrs = ["mongodb_client_container.h"],
148+
data = [
149+
"//src/stirling/source_connectors/socket_tracer/testing/containers/mongodb:client_image.tar",
150+
],
151+
deps = ["//src/common/testing/test_utils:cc_library"],
152+
)
153+
154+
pl_cc_test_library(
155+
name = "mongodb_container",
156+
srcs = [],
157+
hdrs = ["mongodb_container.h"],
158+
data = [
159+
"//src/stirling/source_connectors/socket_tracer/testing/containers:mongodb_image.tar",
160+
],
161+
deps = ["//src/common/testing/test_utils:cc_library"],
162+
)
163+
144164
pl_cc_test_library(
145165
name = "mysql_container",
146166
srcs = [],

‎src/stirling/source_connectors/socket_tracer/testing/container_images/mongodb_client_container.h‎

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,45 @@
1+
/*
2+
* Copyright 2018- The Pixie Authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*
16+
* SPDX-License-Identifier: Apache-2.0
17+
*/
18+
19+
#pragma once
20+
21+
#include <string>
22+
23+
#include "src/common/testing/test_environment.h"
24+
#include "src/common/testing/test_utils/container_runner.h"
25+
26+
namespace px {
27+
namespace stirling {
28+
namespace testing {
29+
30+
class MongoDBClientContainer : public ContainerRunner {
31+
public:
32+
MongoDBClientContainer()
33+
: ContainerRunner(::px::testing::BazelRunfilePath(kBazelImageTar), kContainerNamePrefix,
34+
kReadyMessage) {}
35+
36+
private:
37+
static constexpr std::string_view kBazelImageTar =
38+
"src/stirling/source_connectors/socket_tracer/testing/containers/mongodb/client_image.tar";
39+
static constexpr std::string_view kContainerNamePrefix = "mongodb_client";
40+
static constexpr std::string_view kReadyMessage = "Starting MongoDB client";
41+
};
42+
43+
} // namespace testing
44+
} // namespace stirling
45+
} // namespace px

‎src/stirling/source_connectors/socket_tracer/testing/container_images/mongodb_container.h‎

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,45 @@
1+
/*
2+
* Copyright 2018- The Pixie Authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*
16+
* SPDX-License-Identifier: Apache-2.0
17+
*/
18+
19+
#pragma once
20+
21+
#include <string>
22+
23+
#include "src/common/testing/test_environment.h"
24+
#include "src/common/testing/test_utils/container_runner.h"
25+
26+
namespace px {
27+
namespace stirling {
28+
namespace testing {
29+
30+
class MongoDBContainer : public ContainerRunner {
31+
public:
32+
MongoDBContainer()
33+
: ContainerRunner(::px::testing::BazelRunfilePath(kBazelImageTar), kContainerNamePrefix,
34+
kReadyMessage) {}
35+
36+
private:
37+
static constexpr std::string_view kBazelImageTar =
38+
"src/stirling/source_connectors/socket_tracer/testing/containers/mongodb_image.tar";
39+
static constexpr std::string_view kContainerNamePrefix = "mongodb_server";
40+
static constexpr std::string_view kReadyMessage = "Waiting for connections";
41+
};
42+
43+
} // namespace testing
44+
} // namespace stirling
45+
} // namespace px

0 commit comments

Comments
 (0)