From 6fef7b6fc8b3dc2c0be4d42a62110adcdb5c598e Mon Sep 17 00:00:00 2001
From: Dale Hamel <dale.hamel@shopify.com>
Date: Mon, 16 Sep 2019 20:13:23 -0400
Subject: [PATCH] Add configurable deadline grace period

This allows for bpftrace to take some extra time to process and print
bpf map data.

For very large maps or in some edge cases the user may want to override this value
---
 pkg/cmd/run.go      | 58 +++++++++++++++++++++++++--------------------
 pkg/tracejob/job.go | 33 +++++++++++++-------------
 2 files changed, 49 insertions(+), 42 deletions(-)

diff --git a/pkg/cmd/run.go b/pkg/cmd/run.go
index 00e42da2..44ecb8fb 100644
--- a/pkg/cmd/run.go
+++ b/pkg/cmd/run.go
@@ -25,8 +25,10 @@ var (
 	ImageNameTag = "quay.io/iovisor/kubectl-trace-bpftrace:latest"
 	// InitImageNameTag represents the default init container image
 	InitImageNameTag = "quay.io/iovisor/kubectl-trace-init:latest"
-	// DefaultDeadline default do not allow traces to run for longer than one hour
+	// DefaultDeadline is the maximum time a tracejob is allowed to run, in seconds
 	DefaultDeadline = 3600
+	// DefaultDeadlineGracePeriod is the maximum time to wait to print a map or histogram, in seconds
+	DefaultDeadlineGracePeriod = 10
 )
 
 var (
@@ -68,14 +70,15 @@ type RunOptions struct {
 	explicitNamespace bool
 
 	// Flags local to this command
-	container      string
-	eval           string
-	program        string
-	serviceAccount string
-	imageName      string
-	initImageName  string
-	fetchHeaders   bool
-	deadline       int64
+	container           string
+	eval                string
+	program             string
+	serviceAccount      string
+	imageName           string
+	initImageName       string
+	fetchHeaders        bool
+	deadline            int64
+	deadlineGracePeriod int64
 
 	resourceArg string
 	attach      bool
@@ -91,10 +94,11 @@ func NewRunOptions(streams genericclioptions.IOStreams) *RunOptions {
 	return &RunOptions{
 		IOStreams: streams,
 
-		serviceAccount: "default",
-		imageName:      ImageNameTag,
-		initImageName:  InitImageNameTag,
-		deadline:       int64(DefaultDeadline),
+		serviceAccount:      "default",
+		imageName:           ImageNameTag,
+		initImageName:       InitImageNameTag,
+		deadline:            int64(DefaultDeadline),
+		deadlineGracePeriod: int64(DefaultDeadlineGracePeriod),
 	}
 }
 
@@ -132,6 +136,7 @@ func NewRunCommand(factory factory.Factory, streams genericclioptions.IOStreams)
 	cmd.Flags().StringVar(&o.initImageName, "init-imagename", o.initImageName, "Custom image for the init container responsible to fetch and prepare linux headers")
 	cmd.Flags().BoolVar(&o.fetchHeaders, "fetch-headers", o.fetchHeaders, "Whether to fetch linux headers or not")
 	cmd.Flags().Int64Var(&o.deadline, "deadline", o.deadline, "Maximum time to allow trace to run in seconds")
+	cmd.Flags().Int64Var(&o.deadline, "deadline-grace-period", o.deadlineGracePeriod, "Maximum wait time to print maps or histograms after deadline, in seconds")
 
 	return cmd
 }
@@ -294,19 +299,20 @@ func (o *RunOptions) Run() error {
 	}
 
 	tj := tracejob.TraceJob{
-		Name:             fmt.Sprintf("%s%s", meta.ObjectNamePrefix, string(juid)),
-		Namespace:        o.namespace,
-		ServiceAccount:   o.serviceAccount,
-		ID:               juid,
-		Hostname:         o.nodeName,
-		Program:          o.program,
-		PodUID:           o.podUID,
-		ContainerName:    o.container,
-		IsPod:            o.isPod,
-		ImageNameTag:     o.imageName,
-		InitImageNameTag: o.initImageName,
-		FetchHeaders:     o.fetchHeaders,
-		Deadline:         o.deadline,
+		Name:                fmt.Sprintf("%s%s", meta.ObjectNamePrefix, string(juid)),
+		Namespace:           o.namespace,
+		ServiceAccount:      o.serviceAccount,
+		ID:                  juid,
+		Hostname:            o.nodeName,
+		Program:             o.program,
+		PodUID:              o.podUID,
+		ContainerName:       o.container,
+		IsPod:               o.isPod,
+		ImageNameTag:        o.imageName,
+		InitImageNameTag:    o.initImageName,
+		FetchHeaders:        o.fetchHeaders,
+		Deadline:            o.deadline,
+		DeadlineGracePeriod: o.deadlineGracePeriod,
 	}
 
 	job, err := tc.CreateJob(tj)
diff --git a/pkg/tracejob/job.go b/pkg/tracejob/job.go
index 01a1950e..4ab349a3 100644
--- a/pkg/tracejob/job.go
+++ b/pkg/tracejob/job.go
@@ -23,21 +23,22 @@ type TraceJobClient struct {
 
 // TraceJob is a container of info needed to create the job responsible for tracing.
 type TraceJob struct {
-	Name             string
-	ID               types.UID
-	Namespace        string
-	ServiceAccount   string
-	Hostname         string
-	Program          string
-	PodUID           string
-	ContainerName    string
-	IsPod            bool
-	ImageNameTag     string
-	InitImageNameTag string
-	FetchHeaders     bool
-	Deadline         int64
-	StartTime        metav1.Time
-	Status           TraceJobStatus
+	Name                string
+	ID                  types.UID
+	Namespace           string
+	ServiceAccount      string
+	Hostname            string
+	Program             string
+	PodUID              string
+	ContainerName       string
+	IsPod               bool
+	ImageNameTag        string
+	InitImageNameTag    string
+	FetchHeaders        bool
+	Deadline            int64
+	DeadlineGracePeriod int64
+	StartTime           metav1.Time
+	Status              TraceJobStatus
 }
 
 // WithOutStream setup a file stream to output trace job operation information
@@ -306,7 +307,7 @@ func (t *TraceJobClient) CreateJob(nj TraceJob) (*batchv1.Job, error) {
 										Command: []string{
 											"/bin/bash",
 											"-c",
-											"kill -SIGINT $(pidof bpftrace) && sleep 10",
+											fmt.Sprintf("kill -SIGINT $(pidof bpftrace) && sleep %i", nj.DeadlineGracePeriod),
 										},
 									},
 								},