Skip to content

Commit 598ea47

Browse files
dcantahdcantah
authored
Add IO tracking option for job objects (#1459)
* Add IO tracking option for job objects HCS enables this to get more in depth IO stats for the silo of the container. I'd swapped hostprocess containers to querying for these stats but without the prerequisite of actually enabling them :) This change adds a new option on jobobject.Options{} to enable this functionality and adds a new test to ensure we can actually call StorageStats now. Signed-off-by: Daniel Canter <[email protected]>
1 parent 84e0f9d commit 598ea47

3 files changed

Lines changed: 123 additions & 3 deletions

File tree

‎internal/jobcontainers/jobcontainer.go‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -118,8 +118,9 @@ func Create(ctx context.Context, id string, s *specs.Spec) (_ cow.Container, _ *
118118

119119
// Create the job object all processes will run in.
120120
options := &jobobject.Options{
121-
Name: fmt.Sprintf(jobContainerNameFmt, id),
122-
Notifications: true,
121+
Name: fmt.Sprintf(jobContainerNameFmt, id),
122+
Notifications: true,
123+
EnableIOTracking: true,
123124
}
124125
container.job, err = jobobject.Create(ctx, options)
125126
if err != nil {

‎internal/jobobject/jobobject.go‎

Lines changed: 40 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -72,6 +72,9 @@ type Options struct {
7272
// `Silo` specifies to promote the job to a silo. This additionally sets the flag
7373
// JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE as it is required for the upgrade to complete.
7474
Silo bool
75+
// `IOTracking` enables tracking I/O statistics on the job object. More specifically this
76+
// calls SetInformationJobObject with the JobObjectIoAttribution class.
77+
EnableIOTracking bool
7578
}
7679

7780
// Create creates a job object.
@@ -138,6 +141,12 @@ func Create(ctx context.Context, options *Options) (_ *JobObject, err error) {
138141
job.mq = mq
139142
}
140143

144+
if options.EnableIOTracking {
145+
if err := enableIOTracking(jobHandle); err != nil {
146+
return nil, err
147+
}
148+
}
149+
141150
if options.Silo {
142151
// This is a required setting for upgrading to a silo.
143152
if err := job.SetTerminateOnLastHandleClose(); err != nil {
@@ -433,7 +442,9 @@ func (job *JobObject) QueryProcessorStats() (*winapi.JOBOBJECT_BASIC_ACCOUNTING_
433442
return &info, nil
434443
}
435444

436-
// QueryStorageStats gets the storage (I/O) stats for the job object.
445+
// QueryStorageStats gets the storage (I/O) stats for the job object. This call will error
446+
// if either `EnableIOTracking` wasn't set to true on creation of the job, or SetIOTracking()
447+
// hasn't been called since creation of the job.
437448
func (job *JobObject) QueryStorageStats() (*winapi.JOBOBJECT_IO_ATTRIBUTION_INFORMATION, error) {
438449
job.handleLock.RLock()
439450
defer job.handleLock.RUnlock()
@@ -628,3 +639,31 @@ func (job *JobObject) QueryPrivateWorkingSet() (uint64, error) {
628639

629640
return jobWorkingSetSize, nil
630641
}
642+
643+
// SetIOTracking enables IO tracking for processes in the job object.
644+
// This enables use of the QueryStorageStats method.
645+
func (job *JobObject) SetIOTracking() error {
646+
job.handleLock.RLock()
647+
defer job.handleLock.RUnlock()
648+
649+
if job.handle == 0 {
650+
return ErrAlreadyClosed
651+
}
652+
653+
return enableIOTracking(job.handle)
654+
}
655+
656+
func enableIOTracking(job windows.Handle) error {
657+
info := winapi.JOBOBJECT_IO_ATTRIBUTION_INFORMATION{
658+
ControlFlags: winapi.JOBOBJECT_IO_ATTRIBUTION_CONTROL_ENABLE,
659+
}
660+
if _, err := windows.SetInformationJobObject(
661+
job,
662+
winapi.JobObjectIoAttribution,
663+
uintptr(unsafe.Pointer(&info)),
664+
uint32(unsafe.Sizeof(info)),
665+
); err != nil {
666+
return fmt.Errorf("failed to enable IO tracking on job object: %w", err)
667+
}
668+
return nil
669+
}

‎internal/jobobject/jobobject_test.go‎

Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -65,6 +65,86 @@ func TestSiloCreateAndOpen(t *testing.T) {
6565
}
6666
}
6767

68+
func TestJobStats(t *testing.T) {
69+
var (
70+
ctx = context.Background()
71+
options = &Options{
72+
Name: "test",
73+
Silo: true,
74+
EnableIOTracking: true,
75+
}
76+
)
77+
job, err := Create(ctx, options)
78+
if err != nil {
79+
t.Fatal(err)
80+
}
81+
defer job.Close()
82+
83+
_, err = createProcsAndAssign(1, job)
84+
if err != nil {
85+
t.Fatal(err)
86+
}
87+
88+
_, err = job.QueryMemoryStats()
89+
if err != nil {
90+
t.Fatal(err)
91+
}
92+
93+
_, err = job.QueryProcessorStats()
94+
if err != nil {
95+
t.Fatal(err)
96+
}
97+
98+
_, err = job.QueryStorageStats()
99+
if err != nil {
100+
t.Fatal(err)
101+
}
102+
103+
if err := job.Terminate(1); err != nil {
104+
t.Fatal(err)
105+
}
106+
}
107+
108+
func TestIOTracking(t *testing.T) {
109+
var (
110+
ctx = context.Background()
111+
options = &Options{
112+
Name: "test",
113+
Silo: true,
114+
}
115+
)
116+
job, err := Create(ctx, options)
117+
if err != nil {
118+
t.Fatal(err)
119+
}
120+
defer job.Close()
121+
122+
_, err = createProcsAndAssign(1, job)
123+
if err != nil {
124+
t.Fatal(err)
125+
}
126+
127+
_, err = job.QueryStorageStats()
128+
// Element not found is returned if IO tracking isn't enabled.
129+
if err != nil && !errors.Is(err, windows.ERROR_NOT_FOUND) {
130+
t.Fatal(err)
131+
}
132+
133+
// Turn it on and now the call should function.
134+
if err := job.SetIOTracking(); err != nil {
135+
t.Fatal(err)
136+
}
137+
138+
_, err = job.QueryStorageStats()
139+
if err != nil {
140+
t.Fatal(err)
141+
}
142+
143+
if err := job.Terminate(1); err != nil {
144+
t.Fatal(err)
145+
}
146+
}
147+
68148
func createProcsAndAssign(num int, job *JobObject) (_ []*exec.Cmd, err error) {
69149
var procs []*exec.Cmd
70150

0 commit comments

Comments
 (0)