Skip to content

Commit acf43b0

Browse files
phensleyfntlnz
phensley
authored and
fntlnz
committed
feat: Allow patches to be applied to the job spec
1 parent e3f4d2c commit acf43b0

File tree

3 files changed

+188
-7
lines changed

3 files changed

+188
-7
lines changed

‎pkg/cmd/run.go

+27-7
Original file line numberDiff line numberDiff line change
@@ -58,13 +58,15 @@ var (
5858
# Run a bpftrace inline program on a pod container with a custom image for the bpftrace container that will run your program in the cluster
5959
%[1]s trace run pod/nginx nginx -e "tracepoint:syscalls:sys_enter_* { @[probe] = count(); } --imagename=quay.io/custom-bpftrace-image-name"`
6060

61-
runCommand = "run"
62-
usageString = "(POD | TYPE/NAME)"
63-
requiredArgErrString = fmt.Sprintf("%s is a required argument for the %s command", usageString, runCommand)
64-
containerAsArgOrFlagErrString = "specify container inline as argument or via its flag"
65-
bpftraceMissingErrString = "the bpftrace program is mandatory"
66-
bpftraceDoubleErrString = "specify the bpftrace program either via an external file or via a literal string, not both"
67-
bpftraceEmptyErrString = "the bpftrace programm cannot be empty"
61+
runCommand = "run"
62+
usageString = "(POD | TYPE/NAME)"
63+
requiredArgErrString = fmt.Sprintf("%s is a required argument for the %s command", usageString, runCommand)
64+
containerAsArgOrFlagErrString = "specify container inline as argument or via its flag"
65+
bpftraceMissingErrString = "the bpftrace program is mandatory"
66+
bpftraceDoubleErrString = "specify the bpftrace program either via an external file or via a literal string, not both"
67+
bpftraceEmptyErrString = "the bpftrace programm cannot be empty"
68+
bpftracePatchWithoutTypeErrString = "to use --patch you must also specify the --patch-type argument"
69+
bpftracePatchTypeWithoutPatchErrString = "to use --patch-type you must specify the --patch argument"
6870
)
6971

7072
// RunOptions ...
@@ -91,6 +93,9 @@ type RunOptions struct {
9193
podUID string
9294
nodeName string
9395

96+
patch string
97+
patchType string
98+
9499
clientConfig *rest.Config
95100
}
96101

@@ -142,6 +147,8 @@ func NewRunCommand(factory cmdutil.Factory, streams genericclioptions.IOStreams)
142147
cmd.Flags().BoolVar(&o.fetchHeaders, "fetch-headers", o.fetchHeaders, "Whether to fetch linux headers or not")
143148
cmd.Flags().Int64Var(&o.deadline, "deadline", o.deadline, "Maximum time to allow trace to run in seconds")
144149
cmd.Flags().Int64Var(&o.deadlineGracePeriod, "deadline-grace-period", o.deadlineGracePeriod, "Maximum wait time to print maps or histograms after deadline, in seconds")
150+
cmd.Flags().StringVar(&o.patch, "patch", "", "path of YAML or JSON file used to patch the job definition before creation")
151+
cmd.Flags().StringVar(&o.patchType, "patch-type", "", "patch strategy to use: json, merge, or strategic")
145152

146153
return cmd
147154
}
@@ -175,6 +182,17 @@ func (o *RunOptions) Validate(cmd *cobra.Command, args []string) error {
175182
return fmt.Errorf(bpftraceEmptyErrString)
176183
}
177184

185+
havePatch := cmd.Flag("patch").Changed
186+
havePatchType := cmd.Flag("patch-type").Changed
187+
188+
if havePatch && !havePatchType {
189+
return fmt.Errorf(bpftracePatchWithoutTypeErrString)
190+
}
191+
192+
if !havePatch && havePatchType {
193+
return fmt.Errorf(bpftracePatchTypeWithoutPatchErrString)
194+
}
195+
178196
return nil
179197
}
180198

@@ -318,6 +336,8 @@ func (o *RunOptions) Run() error {
318336
FetchHeaders: o.fetchHeaders,
319337
Deadline: o.deadline,
320338
DeadlineGracePeriod: o.deadlineGracePeriod,
339+
Patch: o.patch,
340+
PatchType: o.patchType,
321341
}
322342

323343
job, err := tc.CreateJob(tj)

‎pkg/tracejob/job.go

+86
Original file line numberDiff line numberDiff line change
@@ -2,19 +2,24 @@ package tracejob
22

33
import (
44
"context"
5+
"encoding/json"
56
"fmt"
67
"io"
78
"io/ioutil"
89
"strconv"
910

11+
jsonpatch "github.com/evanphx/json-patch"
1012
"github.com/iovisor/kubectl-trace/pkg/meta"
1113
batchv1 "k8s.io/api/batch/v1"
1214
apiv1 "k8s.io/api/core/v1"
1315
"k8s.io/apimachinery/pkg/api/resource"
1416
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
1517
"k8s.io/apimachinery/pkg/types"
18+
"k8s.io/apimachinery/pkg/util/strategicpatch"
19+
yamlutil "k8s.io/apimachinery/pkg/util/yaml"
1620
batchv1typed "k8s.io/client-go/kubernetes/typed/batch/v1"
1721
corev1typed "k8s.io/client-go/kubernetes/typed/core/v1"
22+
"sigs.k8s.io/yaml"
1823
)
1924

2025
type TraceJobClient struct {
@@ -41,6 +46,8 @@ type TraceJob struct {
4146
DeadlineGracePeriod int64
4247
StartTime *metav1.Time
4348
Status TraceJobStatus
49+
Patch string
50+
PatchType string
4451
}
4552

4653
// WithOutStream setup a file stream to output trace job operation information
@@ -473,6 +480,16 @@ func (t *TraceJobClient) CreateJob(nj TraceJob) (*batchv1.Job, error) {
473480
if _, err := t.ConfigClient.Create(context.Background(), cm, metav1.CreateOptions{}); err != nil {
474481
return nil, err
475482
}
483+
484+
// Optionally patch the job before creating it
485+
if nj.PatchType != "" && nj.Patch != "" {
486+
newJob, err := patchJobFile(job, nj.PatchType, nj.Patch)
487+
if err != nil {
488+
return nil, err
489+
}
490+
job = newJob
491+
}
492+
476493
return t.JobClient.Create(context.Background(), job, metav1.CreateOptions{})
477494
}
478495

@@ -548,3 +565,72 @@ func jobStatus(j batchv1.Job) TraceJobStatus {
548565
}
549566
return TraceJobUnknown
550567
}
568+
569+
var patchTypes = map[string]types.PatchType{
570+
"json": types.JSONPatchType,
571+
"merge": types.MergePatchType,
572+
"strategic": types.StrategicMergePatchType,
573+
}
574+
575+
func patchJobFile(j *batchv1.Job, patchType, patchPath string) (*batchv1.Job, error) {
576+
patchYAML, err := ioutil.ReadFile(patchPath)
577+
if err != nil {
578+
return nil, fmt.Errorf("failed to read patch yaml path %v: %s", patchPath, err)
579+
}
580+
return patchJob(j, patchType, patchYAML)
581+
}
582+
583+
func patchJob(j *batchv1.Job, patchType string, patchBytes []byte) (*batchv1.Job, error) {
584+
var err error
585+
patchJSON := patchBytes
586+
587+
if !json.Valid(patchBytes) {
588+
// Convert YAML to JSON for patching
589+
patchJSON, err = yamlutil.ToJSON(patchBytes)
590+
if err != nil {
591+
return nil, fmt.Errorf("converting patch yaml to json: %s", err)
592+
}
593+
}
594+
595+
jobYAML, err := yaml.Marshal(j)
596+
if err != nil {
597+
return nil, fmt.Errorf("marshal job to yaml: %s", err)
598+
}
599+
600+
jobJSON, err := yamlutil.ToJSON(jobYAML)
601+
if err != nil {
602+
return nil, fmt.Errorf("converting job yaml to json: %s", err)
603+
}
604+
605+
// Patch job JSON
606+
typ := patchTypes[patchType]
607+
switch typ {
608+
case types.JSONPatchType:
609+
raw, err := jsonpatch.DecodePatch(patchJSON)
610+
if err != nil {
611+
return nil, fmt.Errorf("decoding json patch: %s", err)
612+
}
613+
jobJSON, err = raw.Apply(jobJSON)
614+
615+
case types.MergePatchType:
616+
jobJSON, err = jsonpatch.MergePatch(jobJSON, patchJSON)
617+
618+
case types.StrategicMergePatchType:
619+
jobJSON, err = strategicpatch.StrategicMergePatch(jobJSON, patchJSON, batchv1.Job{})
620+
621+
default:
622+
return nil, fmt.Errorf("%v is an invalid patch type", patchType)
623+
}
624+
625+
if err != nil {
626+
return nil, fmt.Errorf("applying %s patch to job: %s", patchType, err)
627+
}
628+
629+
// Unmarshal back to Job object
630+
newJob := batchv1.Job{}
631+
if err = json.Unmarshal(jobJSON, &newJob); err != nil {
632+
return nil, fmt.Errorf("failed to marshal job from patched json: %s", err)
633+
}
634+
635+
return &newJob, nil
636+
}

‎pkg/tracejob/job_test.go

+75
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,75 @@
1+
package tracejob
2+
3+
import (
4+
"reflect"
5+
"testing"
6+
7+
batchv1 "k8s.io/api/batch/v1"
8+
)
9+
10+
type patchTest struct {
11+
patchType string
12+
patch []byte
13+
}
14+
15+
var patchJSON1 = []byte(`
16+
- op: replace
17+
path: "/spec/backOffLimit"
18+
value: 123
19+
- op: add
20+
path: "/spec/template/hostPID"
21+
value: true
22+
- op: remove
23+
path: "/spec/completions"
24+
`)
25+
var patchJSON2 = []byte(`[
26+
{ "op": "replace", "path": "/spec/backOffLimit", "value": 123 },
27+
{ "op": "add", "path": "/spec/template/hostPID", "value": true },
28+
{ "op": "remove", "path": "/spec/completions"}
29+
]`)
30+
31+
var patchMerge = []byte(`
32+
spec:
33+
backoffLimit: 123
34+
completions: null
35+
template:
36+
hostPID: true
37+
`)
38+
39+
var testCases = []patchTest{
40+
{patchType: "json", patch: patchJSON1},
41+
{patchType: "json", patch: patchJSON2},
42+
{patchType: "merge", patch: patchMerge},
43+
{patchType: "strategic", patch: patchMerge},
44+
}
45+
46+
func TestPatchJobJSON(t *testing.T) {
47+
for _, c := range testCases {
48+
job := getJob()
49+
newJob, err := patchJob(job, c.patchType, c.patch)
50+
if err != nil {
51+
t.Error(err)
52+
}
53+
54+
// Update expected value
55+
job.Spec.BackoffLimit = int32Ptr(123)
56+
job.Spec.Template.Spec.HostPID = true
57+
job.Spec.Completions = nil
58+
59+
if reflect.DeepEqual(job, newJob) {
60+
t.Errorf("patch %s job does not match expected", c.patchType)
61+
}
62+
}
63+
}
64+
65+
func getJob() *batchv1.Job {
66+
return &batchv1.Job{
67+
Spec: batchv1.JobSpec{
68+
ActiveDeadlineSeconds: int64Ptr(60),
69+
TTLSecondsAfterFinished: int32Ptr(5),
70+
Parallelism: int32Ptr(1),
71+
Completions: int32Ptr(1),
72+
BackoffLimit: int32Ptr(1),
73+
},
74+
}
75+
}

0 commit comments

Comments
 (0)