Skip to content

Commit 52e003d

Browse files
benkilimnikbenkilimnik
authored
[Kernel scoping 4/5] Add TraceProgram object (#1687)
1 parent 1b65ee8 commit 52e003d

File tree

7 files changed

+450
-25
lines changed

7 files changed

+450
-25
lines changed

‎src/carnot/planner/logical_planner_test.cc‎

Lines changed: 81 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -373,6 +373,85 @@ TEST_F(LogicalPlannerTest, PlanWithExecFuncs) {
373373
EXPECT_OK(plan->ToProto());
374374
}
375375

376+
constexpr char kBPFTraceProgramMaxKernel[] = R"bpftrace(
377+
kprobe:tcp_drop
378+
{
379+
...
380+
}
381+
)bpftrace";
382+
383+
constexpr char kBPFTraceProgramMinKernel[] = R"bpftrace(
384+
tracepoint:skb:kfree_skb
385+
{
386+
...
387+
}
388+
)bpftrace";
389+
390+
constexpr char kTwoTraceProgramsPxl[] = R"pxl(
391+
import pxtrace
392+
import px
393+
394+
before_518_trace_program = pxtrace.TraceProgram(
395+
program="""$0""",
396+
max_kernel='5.18',
397+
)
398+
399+
after_519_trace_program = pxtrace.TraceProgram(
400+
program="""$1""",
401+
min_kernel='5.19',
402+
)
403+
404+
table_name = 'tcp_drop_table'
405+
pxtrace.UpsertTracepoint('tcp_drop_tracer',
406+
table_name,
407+
[before_518_trace_program, after_519_trace_program],
408+
pxtrace.kprobe(),
409+
'10m')
410+
)pxl";
411+
412+
constexpr char kBPFTwoTraceProgramsPb[] = R"proto(
413+
name: "tcp_drop_tracer"
414+
ttl {
415+
seconds: 600
416+
}
417+
programs {
418+
table_name: "tcp_drop_table"
419+
bpftrace {
420+
program: "\nkprobe:tcp_drop\n{\n ...\n}\n"
421+
}
422+
selectors {
423+
selector_type: MAX_KERNEL
424+
value: "5.18"
425+
}
426+
}
427+
programs {
428+
table_name: "tcp_drop_table"
429+
bpftrace {
430+
program: "\ntracepoint:skb:kfree_skb\n{\n ...\n}\n"
431+
}
432+
selectors {
433+
selector_type: MIN_KERNEL
434+
value: "5.19"
435+
}
436+
}
437+
)proto";
438+
439+
TEST_F(LogicalPlannerTest, CompileTwoTracePrograms) {
440+
auto planner = LogicalPlanner::Create(info_).ConsumeValueOrDie();
441+
plannerpb::CompileMutationsRequest req;
442+
req.set_query_str(
443+
absl::Substitute(kTwoTraceProgramsPxl, kBPFTraceProgramMaxKernel, kBPFTraceProgramMinKernel));
444+
*req.mutable_logical_planner_state() =
445+
testutils::CreateTwoPEMsOneKelvinPlannerState(testutils::kHttpEventsSchema);
446+
auto trace_ir_or_s = planner->CompileTrace(req);
447+
ASSERT_OK(trace_ir_or_s);
448+
auto trace_ir = trace_ir_or_s.ConsumeValueOrDie();
449+
plannerpb::CompileMutationsResponse resp;
450+
ASSERT_OK(trace_ir->ToProto(&resp));
451+
ASSERT_EQ(resp.mutations_size(), 1);
452+
EXPECT_THAT(resp.mutations()[0].trace(), EqualsProto(kBPFTwoTraceProgramsPb));
453+
}
454+
376455
constexpr char kSingleProbePxl[] = R"pxl(
377456
import pxtrace
378457
import px
@@ -391,7 +470,7 @@ pxtrace.UpsertTracepoint('http_return',
391470
"5m")
392471
)pxl";
393472

394-
constexpr char kSingleProbeProgramPb[] = R"pxl(
473+
constexpr char kSingleProbeProgramPb[] = R"proto(
395474
name: "http_return"
396475
ttl {
397476
seconds: 300
@@ -435,7 +514,7 @@ programs {
435514
}
436515
}
437516
}
438-
)pxl";
517+
)proto";
439518

440519
TEST_F(LogicalPlannerTest, CompileTrace) {
441520
auto planner = LogicalPlanner::Create(info_).ConsumeValueOrDie();

‎src/carnot/planner/objects/qlobject.h‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,7 @@ enum class QLObjectType {
5555
// General module type.
5656
kModule,
5757
kTraceModule,
58+
kTraceProgram,
5859
kDict,
5960
kTracingVariable,
6061
kProbe,

‎src/carnot/planner/probes/probes.cc‎

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -197,12 +197,17 @@ StatusOr<TracepointDeployment*> MutationsIR::CreateKProbeTracepointDeployment(
197197
return raw;
198198
}
199199

200-
Status TracepointDeployment::AddBPFTrace(const std::string& bpftrace,
201-
const std::string& output_name) {
200+
Status TracepointDeployment::AddBPFTrace(const std::string& bpftrace_str,
201+
const std::string& output_name,
202+
const std::vector<TracepointSelector>& selectors) {
202203
carnot::planner::dynamic_tracing::ir::logical::TracepointDeployment::TracepointProgram
203204
tracepoint_pb;
204-
tracepoint_pb.mutable_bpftrace()->set_program(bpftrace);
205+
tracepoint_pb.mutable_bpftrace()->set_program(bpftrace_str);
206+
// set the output table to write program results to
205207
tracepoint_pb.set_table_name(output_name);
208+
for (const auto& selector : selectors) {
209+
*tracepoint_pb.add_selectors() = selector;
210+
}
206211
tracepoints_.push_back(tracepoint_pb);
207212
return Status::OK();
208213
}

‎src/carnot/planner/probes/probes.h‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,8 @@ namespace carnot {
3636
namespace planner {
3737
namespace compiler {
3838

39+
using TracepointSelector = carnot::planner::dynamic_tracing::ir::logical::TracepointSelector;
40+
3941
class ProbeOutput {
4042
public:
4143
ProbeOutput() = delete;
@@ -192,9 +194,11 @@ class TracepointDeployment {
192194
*
193195
* @param bpftrace_program the program in string format.
194196
* @param output_name the output table to write program results.
197+
* @param selectors the selectors to use for the program.
195198
* @return Status
196199
*/
197-
Status AddBPFTrace(const std::string& bpftrace_program, const std::string& output_name);
200+
Status AddBPFTrace(const std::string& bpftrace_str, const std::string& output_name,
201+
const std::vector<TracepointSelector>& selectors);
198202

199203
std::string name() const { return name_; }
200204

‎src/carnot/planner/probes/probes_test.cc‎

Lines changed: 213 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -648,6 +648,219 @@ TEST_F(ProbeCompilerTest, parse_bpftrace) {
648648
testing::proto::EqualsProto(absl::Substitute(kBPFTraceProgramPb, literal_bpf_trace)));
649649
}
650650

651+
constexpr char kBPFTraceProgramMaxKernel[] = R"bpftrace(
652+
kprobe:tcp_drop
653+
{
654+
...
655+
}
656+
)bpftrace";
657+
658+
constexpr char kBPFTraceProgramMinKernel[] = R"bpftrace(
659+
tracepoint:skb:kfree_skb
660+
{
661+
...
662+
}
663+
)bpftrace";
664+
665+
// Test that we can compile/parse a single TraceProgram object with a valid selector
666+
constexpr char kBPFSingleTraceProgramObjectPxl[] = R"pxl(
667+
import pxtrace
668+
import px
669+
670+
after_519_trace_program = pxtrace.TraceProgram(
671+
program="""$0""",
672+
min_kernel='5.19',
673+
)
674+
675+
table_name = 'tcp_drop_table'
676+
pxtrace.UpsertTracepoint('tcp_drop_tracer',
677+
table_name,
678+
after_519_trace_program,
679+
pxtrace.kprobe(),
680+
'10m')
681+
)pxl";
682+
683+
constexpr char kBPFSingleTraceProgramObjectPb[] = R"proto(
684+
name: "tcp_drop_tracer"
685+
ttl {
686+
seconds: 600
687+
}
688+
programs {
689+
table_name: "tcp_drop_table"
690+
bpftrace {
691+
program: "$0"
692+
}
693+
selectors {
694+
selector_type: MIN_KERNEL
695+
value: "5.19"
696+
}
697+
}
698+
)proto";
699+
700+
TEST_F(ProbeCompilerTest, parse_single_bpftrace_program_object) {
701+
ASSERT_OK_AND_ASSIGN(auto probe_ir,
702+
CompileProbeScript(absl::Substitute(kBPFSingleTraceProgramObjectPxl,
703+
kBPFTraceProgramMinKernel)));
704+
plannerpb::CompileMutationsResponse pb;
705+
EXPECT_OK(probe_ir->ToProto(&pb));
706+
ASSERT_EQ(pb.mutations_size(), 1);
707+
708+
std::string literal_bpf_trace_min = kBPFTraceProgramMinKernel;
709+
literal_bpf_trace_min = std::regex_replace(literal_bpf_trace_min, std::regex("\n"), "\\n");
710+
711+
EXPECT_THAT(pb.mutations()[0].trace(),
712+
testing::proto::EqualsProto(
713+
absl::Substitute(kBPFSingleTraceProgramObjectPb, literal_bpf_trace_min)));
714+
}
715+
716+
// Test that we can compile a list of TraceProgram objects with valid selectors
717+
constexpr char kBPFTraceProgramObjectsPxl[] = R"pxl(
718+
import pxtrace
719+
import px
720+
721+
before_518_trace_program = pxtrace.TraceProgram(
722+
program="""$0""",
723+
max_kernel='5.18',
724+
)
725+
726+
after_519_trace_program = pxtrace.TraceProgram(
727+
program="""$1""",
728+
min_kernel='5.19',
729+
)
730+
731+
table_name = 'tcp_drop_table'
732+
pxtrace.UpsertTracepoint('tcp_drop_tracer',
733+
table_name,
734+
[before_518_trace_program, after_519_trace_program],
735+
pxtrace.kprobe(),
736+
'10m')
737+
)pxl";
738+
739+
constexpr char kBPFTraceProgramObjectsPb[] = R"proto(
740+
name: "tcp_drop_tracer"
741+
ttl {
742+
seconds: 600
743+
}
744+
programs {
745+
table_name: "tcp_drop_table"
746+
bpftrace {
747+
program: "$0"
748+
}
749+
selectors {
750+
selector_type: MAX_KERNEL
751+
value: "5.18"
752+
}
753+
}
754+
programs {
755+
table_name: "tcp_drop_table"
756+
bpftrace {
757+
program: "$1"
758+
}
759+
selectors {
760+
selector_type: MIN_KERNEL
761+
value: "5.19"
762+
}
763+
}
764+
)proto";
765+
766+
TEST_F(ProbeCompilerTest, parse_multiple_bpftrace_program_objects) {
767+
ASSERT_OK_AND_ASSIGN(auto probe_ir, CompileProbeScript(absl::Substitute(
768+
kBPFTraceProgramObjectsPxl, kBPFTraceProgramMinKernel,
769+
kBPFTraceProgramMaxKernel)));
770+
plannerpb::CompileMutationsResponse pb;
771+
EXPECT_OK(probe_ir->ToProto(&pb));
772+
ASSERT_EQ(pb.mutations_size(), 1);
773+
774+
std::string literal_bpf_trace_min = kBPFTraceProgramMinKernel;
775+
literal_bpf_trace_min = std::regex_replace(literal_bpf_trace_min, std::regex("\n"), "\\n");
776+
777+
std::string literal_bpf_trace_max = kBPFTraceProgramMaxKernel;
778+
literal_bpf_trace_max = std::regex_replace(literal_bpf_trace_max, std::regex("\n"), "\\n");
779+
780+
EXPECT_THAT(pb.mutations()[0].trace(),
781+
testing::proto::EqualsProto(absl::Substitute(
782+
kBPFTraceProgramObjectsPb, literal_bpf_trace_min, literal_bpf_trace_max)));
783+
}
784+
785+
// Test that passing an unsupported selector type to TraceProgram throws a compiler error
786+
constexpr char kBPFUnsupportedTraceProgramObjectSelectorPxl[] = R"pxl(
787+
import pxtrace
788+
import px
789+
790+
after_519_trace_program = pxtrace.TraceProgram(
791+
program="""$0""",
792+
min_kernel='5.19',
793+
my_unsupported_selector='12345',
794+
)
795+
796+
table_name = 'tcp_drop_table'
797+
pxtrace.UpsertTracepoint('tcp_drop_tracer',
798+
table_name,
799+
after_519_trace_program,
800+
pxtrace.kprobe(),
801+
'10m')
802+
)pxl";
803+
804+
TEST_F(ProbeCompilerTest, parse_unsupported_selector_in_trace_program_object) {
805+
auto probe_ir_or_s = CompileProbeScript(kBPFUnsupportedTraceProgramObjectSelectorPxl);
806+
ASSERT_NOT_OK(probe_ir_or_s);
807+
EXPECT_THAT(
808+
probe_ir_or_s.status(),
809+
HasCompilerError("Unsupported selector argument provided \'my_unsupported_selector\'"));
810+
}
811+
812+
// Test that an invalid selector value throws a compiler error (currently needs to be a string)
813+
constexpr char kBPFInvalidTraceProgramObjectSelectorPxl[] = R"pxl(
814+
import pxtrace
815+
import px
816+
817+
after_519_trace_program = pxtrace.TraceProgram(
818+
program="""$0""",
819+
min_kernel='5.19',
820+
max_kernel=None,
821+
)
822+
823+
table_name = 'tcp_drop_table'
824+
pxtrace.UpsertTracepoint('tcp_drop_tracer',
825+
table_name,
826+
after_519_trace_program,
827+
pxtrace.kprobe(),
828+
'10m')
829+
)pxl";
830+
831+
TEST_F(ProbeCompilerTest, parse_invalid_trace_program_object) {
832+
auto probe_ir_or_s = CompileProbeScript(kBPFInvalidTraceProgramObjectSelectorPxl);
833+
ASSERT_NOT_OK(probe_ir_or_s);
834+
EXPECT_THAT(probe_ir_or_s.status(),
835+
HasCompilerError("Expected \'String\' in arg \'max_kernel\', got \'none\'"));
836+
}
837+
838+
// Test that an empty selector value throws a compiler error
839+
constexpr char kBPFEmptyTraceProgramObjectSelectorPxl[] = R"pxl(
840+
import pxtrace
841+
import px
842+
843+
after_519_trace_program = pxtrace.TraceProgram(
844+
program="""$0""",
845+
min_kernel='5.19',
846+
max_kernel='',
847+
)
848+
849+
table_name = 'tcp_drop_table'
850+
pxtrace.UpsertTracepoint('tcp_drop_tracer',
851+
table_name,
852+
after_519_trace_program,
853+
pxtrace.kprobe(),
854+
'10m')
855+
)pxl";
856+
857+
TEST_F(ProbeCompilerTest, parse_empty_trace_program_object) {
858+
auto probe_ir_or_s = CompileProbeScript(kBPFEmptyTraceProgramObjectSelectorPxl);
859+
ASSERT_NOT_OK(probe_ir_or_s);
860+
EXPECT_THAT(probe_ir_or_s.status(),
861+
HasCompilerError("Empty selector value provided for \'max_kernel\'"));
862+
}
863+
651864
constexpr char kConfigChangePxl[] = R"pxl(
652865
import pxconfig
653866
import px

0 commit comments

Comments
 (0)