Skip to content

Commit daf54c5

Browse files
benkilimnikbenkilimnik
committed
Add TraceProgram object, rebase main
Signed-off-by: Benjamin Kilimnik <bkilimnik@pixielabs.ai>
1 parent f974c70 commit daf54c5

7 files changed

Lines changed: 458 additions & 25 deletions

File tree

‎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: 222 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -648,6 +648,228 @@ 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 =
710+
std::regex_replace(literal_bpf_trace_min, std::regex(R"(\\\n)"), R"(\\\\n)");
711+
literal_bpf_trace_min = std::regex_replace(literal_bpf_trace_min, std::regex("\n"), "\\n");
712+
literal_bpf_trace_min = std::regex_replace(literal_bpf_trace_min, std::regex("\""), "\\\"");
713+
714+
EXPECT_THAT(pb.mutations()[0].trace(),
715+
testing::proto::EqualsProto(
716+
absl::Substitute(kBPFSingleTraceProgramObjectPb, literal_bpf_trace_min)));
717+
}
718+
719+
// Test that we can compile a list of TraceProgram objects with valid selectors
720+
constexpr char kBPFTraceProgramObjectsPxl[] = R"pxl(
721+
import pxtrace
722+
import px
723+
724+
before_518_trace_program = pxtrace.TraceProgram(
725+
program="""$0""",
726+
max_kernel='5.18',
727+
)
728+
729+
after_519_trace_program = pxtrace.TraceProgram(
730+
program="""$1""",
731+
min_kernel='5.19',
732+
)
733+
734+
table_name = 'tcp_drop_table'
735+
pxtrace.UpsertTracepoint('tcp_drop_tracer',
736+
table_name,
737+
[before_518_trace_program, after_519_trace_program],
738+
pxtrace.kprobe(),
739+
'10m')
740+
)pxl";
741+
742+
constexpr char kBPFTraceProgramObjectsPb[] = R"proto(
743+
name: "tcp_drop_tracer"
744+
ttl {
745+
seconds: 600
746+
}
747+
programs {
748+
table_name: "tcp_drop_table"
749+
bpftrace {
750+
program: "$0"
751+
}
752+
selectors {
753+
selector_type: MAX_KERNEL
754+
value: "5.18"
755+
}
756+
}
757+
programs {
758+
table_name: "tcp_drop_table"
759+
bpftrace {
760+
program: "$1"
761+
}
762+
selectors {
763+
selector_type: MIN_KERNEL
764+
value: "5.19"
765+
}
766+
}
767+
)proto";
768+
769+
TEST_F(ProbeCompilerTest, parse_multiple_bpftrace_program_objects) {
770+
ASSERT_OK_AND_ASSIGN(auto probe_ir, CompileProbeScript(absl::Substitute(
771+
kBPFTraceProgramObjectsPxl, kBPFTraceProgramMinKernel,
772+
kBPFTraceProgramMaxKernel)));
773+
plannerpb::CompileMutationsResponse pb;
774+
EXPECT_OK(probe_ir->ToProto(&pb));
775+
ASSERT_EQ(pb.mutations_size(), 1);
776+
777+
std::string literal_bpf_trace_min = kBPFTraceProgramMinKernel;
778+
literal_bpf_trace_min =
779+
std::regex_replace(literal_bpf_trace_min, std::regex(R"(\\\n)"), R"(\\\\n)");
780+
literal_bpf_trace_min = std::regex_replace(literal_bpf_trace_min, std::regex("\n"), "\\n");
781+
literal_bpf_trace_min = std::regex_replace(literal_bpf_trace_min, std::regex("\""), "\\\"");
782+
783+
std::string literal_bpf_trace_max = kBPFTraceProgramMaxKernel;
784+
literal_bpf_trace_max =
785+
std::regex_replace(literal_bpf_trace_max, std::regex(R"(\\\n)"), R"(\\\\n)");
786+
literal_bpf_trace_max = std::regex_replace(literal_bpf_trace_max, std::regex("\n"), "\\n");
787+
literal_bpf_trace_max = std::regex_replace(literal_bpf_trace_max, std::regex("\""), "\\\"");
788+
789+
EXPECT_THAT(pb.mutations()[0].trace(),
790+
testing::proto::EqualsProto(absl::Substitute(
791+
kBPFTraceProgramObjectsPb, literal_bpf_trace_min, literal_bpf_trace_max)));
792+
}
793+
794+
// Test that passing an unsupported selector type to TraceProgram throws a compiler error
795+
constexpr char kBPFUnsupportedTraceProgramObjectSelectorPxl[] = R"pxl(
796+
import pxtrace
797+
import px
798+
799+
after_519_trace_program = pxtrace.TraceProgram(
800+
program="""$0""",
801+
min_kernel='5.19',
802+
my_unsupported_selector='12345',
803+
)
804+
805+
table_name = 'tcp_drop_table'
806+
pxtrace.UpsertTracepoint('tcp_drop_tracer',
807+
table_name,
808+
after_519_trace_program,
809+
pxtrace.kprobe(),
810+
'10m')
811+
)pxl";
812+
813+
TEST_F(ProbeCompilerTest, parse_unsupported_selector_in_trace_program_object) {
814+
auto probe_ir_or_s = CompileProbeScript(kBPFUnsupportedTraceProgramObjectSelectorPxl);
815+
ASSERT_NOT_OK(probe_ir_or_s);
816+
EXPECT_THAT(
817+
probe_ir_or_s.status(),
818+
HasCompilerError("Unsupported selector argument provided \'my_unsupported_selector\'"));
819+
}
820+
821+
// Test that an invalid selector value throws a compiler error (currently needs to be a string)
822+
constexpr char kBPFInvalidTraceProgramObjectSelectorPxl[] = R"pxl(
823+
import pxtrace
824+
import px
825+
826+
after_519_trace_program = pxtrace.TraceProgram(
827+
program="""$0""",
828+
min_kernel='5.19',
829+
max_kernel=None,
830+
)
831+
832+
table_name = 'tcp_drop_table'
833+
pxtrace.UpsertTracepoint('tcp_drop_tracer',
834+
table_name,
835+
after_519_trace_program,
836+
pxtrace.kprobe(),
837+
'10m')
838+
)pxl";
839+
840+
TEST_F(ProbeCompilerTest, parse_invalid_trace_program_object) {
841+
auto probe_ir_or_s = CompileProbeScript(kBPFInvalidTraceProgramObjectSelectorPxl);
842+
ASSERT_NOT_OK(probe_ir_or_s);
843+
EXPECT_THAT(probe_ir_or_s.status(),
844+
HasCompilerError("Expected \'String\' in arg \'max_kernel\', got \'none\'"));
845+
}
846+
847+
// Test that an empty selector value throws a compiler error
848+
constexpr char kBPFEmptyTraceProgramObjectSelectorPxl[] = R"pxl(
849+
import pxtrace
850+
import px
851+
852+
after_519_trace_program = pxtrace.TraceProgram(
853+
program="""$0""",
854+
min_kernel='5.19',
855+
max_kernel='',
856+
)
857+
858+
table_name = 'tcp_drop_table'
859+
pxtrace.UpsertTracepoint('tcp_drop_tracer',
860+
table_name,
861+
after_519_trace_program,
862+
pxtrace.kprobe(),
863+
'10m')
864+
)pxl";
865+
866+
TEST_F(ProbeCompilerTest, parse_empty_trace_program_object) {
867+
auto probe_ir_or_s = CompileProbeScript(kBPFEmptyTraceProgramObjectSelectorPxl);
868+
ASSERT_NOT_OK(probe_ir_or_s);
869+
EXPECT_THAT(probe_ir_or_s.status(),
870+
HasCompilerError("Empty selector value provided for \'max_kernel\'"));
871+
}
872+
651873
constexpr char kConfigChangePxl[] = R"pxl(
652874
import pxconfig
653875
import px

0 commit comments

Comments
 (0)