limits.go

Documentation: github.com/Microsoft/hcsshim/internal/jobobject

     1  //go:build windows
     2  
     3  package jobobject
     4  
     5  import (
     6  	"errors"
     7  	"fmt"
     8  	"unsafe"
     9  
    10  	"github.com/Microsoft/hcsshim/internal/winapi"
    11  	"golang.org/x/sys/windows"
    12  )
    13  
    14  const (
    15  	memoryLimitMax uint64 = 0xffffffffffffffff
    16  )
    17  
    18  func isFlagSet(flag, controlFlags uint32) bool {
    19  	return (flag & controlFlags) == flag
    20  }
    21  
    22  // SetResourceLimits sets resource limits on the job object (cpu, memory, storage).
    23  func (job *JobObject) SetResourceLimits(limits *JobLimits) error {
    24  	// Go through and check what limits were specified and apply them to the job.
    25  	if limits.MemoryLimitInBytes != 0 {
    26  		if err := job.SetMemoryLimit(limits.MemoryLimitInBytes); err != nil {
    27  			return fmt.Errorf("failed to set job object memory limit: %w", err)
    28  		}
    29  	}
    30  
    31  	if limits.CPULimit != 0 {
    32  		if err := job.SetCPULimit(RateBased, limits.CPULimit); err != nil {
    33  			return fmt.Errorf("failed to set job object cpu limit: %w", err)
    34  		}
    35  	} else if limits.CPUWeight != 0 {
    36  		if err := job.SetCPULimit(WeightBased, limits.CPUWeight); err != nil {
    37  			return fmt.Errorf("failed to set job object cpu limit: %w", err)
    38  		}
    39  	}
    40  
    41  	if limits.MaxBandwidth != 0 || limits.MaxIOPS != 0 {
    42  		if err := job.SetIOLimit(limits.MaxBandwidth, limits.MaxIOPS); err != nil {
    43  			return fmt.Errorf("failed to set io limit on job object: %w", err)
    44  		}
    45  	}
    46  	return nil
    47  }
    48  
    49  // SetTerminateOnLastHandleClose sets the job object flag that specifies that the job should terminate
    50  // all processes in the job on the last open handle being closed.
    51  func (job *JobObject) SetTerminateOnLastHandleClose() error {
    52  	info, err := job.getExtendedInformation()
    53  	if err != nil {
    54  		return err
    55  	}
    56  	info.BasicLimitInformation.LimitFlags |= windows.JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE
    57  	return job.setExtendedInformation(info)
    58  }
    59  
    60  // SetMemoryLimit sets the memory limit of the job object based on the given `memoryLimitInBytes`.
    61  func (job *JobObject) SetMemoryLimit(memoryLimitInBytes uint64) error {
    62  	if memoryLimitInBytes >= memoryLimitMax {
    63  		return errors.New("memory limit specified exceeds the max size")
    64  	}
    65  
    66  	info, err := job.getExtendedInformation()
    67  	if err != nil {
    68  		return err
    69  	}
    70  
    71  	info.JobMemoryLimit = uintptr(memoryLimitInBytes)
    72  	info.BasicLimitInformation.LimitFlags |= windows.JOB_OBJECT_LIMIT_JOB_MEMORY
    73  	return job.setExtendedInformation(info)
    74  }
    75  
    76  // GetMemoryLimit gets the memory limit in bytes of the job object.
    77  func (job *JobObject) GetMemoryLimit() (uint64, error) {
    78  	info, err := job.getExtendedInformation()
    79  	if err != nil {
    80  		return 0, err
    81  	}
    82  	return uint64(info.JobMemoryLimit), nil
    83  }
    84  
    85  // SetCPULimit sets the CPU limit depending on the specified `CPURateControlType` to
    86  // `rateControlValue` for the job object.
    87  func (job *JobObject) SetCPULimit(rateControlType CPURateControlType, rateControlValue uint32) error {
    88  	cpuInfo, err := job.getCPURateControlInformation()
    89  	if err != nil {
    90  		return err
    91  	}
    92  	switch rateControlType {
    93  	case WeightBased:
    94  		if rateControlValue < cpuWeightMin || rateControlValue > cpuWeightMax {
    95  			return fmt.Errorf("processor weight value of `%d` is invalid", rateControlValue)
    96  		}
    97  		cpuInfo.ControlFlags |= winapi.JOB_OBJECT_CPU_RATE_CONTROL_ENABLE | winapi.JOB_OBJECT_CPU_RATE_CONTROL_WEIGHT_BASED
    98  		cpuInfo.Value = rateControlValue
    99  	case RateBased:
   100  		if rateControlValue < cpuLimitMin || rateControlValue > cpuLimitMax {
   101  			return fmt.Errorf("processor rate of `%d` is invalid", rateControlValue)
   102  		}
   103  		cpuInfo.ControlFlags |= winapi.JOB_OBJECT_CPU_RATE_CONTROL_ENABLE | winapi.JOB_OBJECT_CPU_RATE_CONTROL_HARD_CAP
   104  		cpuInfo.Value = rateControlValue
   105  	default:
   106  		return errors.New("invalid job object cpu rate control type")
   107  	}
   108  	return job.setCPURateControlInfo(cpuInfo)
   109  }
   110  
   111  // GetCPULimit gets the cpu limits for the job object.
   112  // `rateControlType` is used to indicate what type of cpu limit to query for.
   113  func (job *JobObject) GetCPULimit(rateControlType CPURateControlType) (uint32, error) {
   114  	info, err := job.getCPURateControlInformation()
   115  	if err != nil {
   116  		return 0, err
   117  	}
   118  
   119  	if !isFlagSet(winapi.JOB_OBJECT_CPU_RATE_CONTROL_ENABLE, info.ControlFlags) {
   120  		return 0, errors.New("the job does not have cpu rate control enabled")
   121  	}
   122  
   123  	switch rateControlType {
   124  	case WeightBased:
   125  		if !isFlagSet(winapi.JOB_OBJECT_CPU_RATE_CONTROL_WEIGHT_BASED, info.ControlFlags) {
   126  			return 0, errors.New("cannot get cpu weight for job object without cpu weight option set")
   127  		}
   128  	case RateBased:
   129  		if !isFlagSet(winapi.JOB_OBJECT_CPU_RATE_CONTROL_HARD_CAP, info.ControlFlags) {
   130  			return 0, errors.New("cannot get cpu rate hard cap for job object without cpu rate hard cap option set")
   131  		}
   132  	default:
   133  		return 0, errors.New("invalid job object cpu rate control type")
   134  	}
   135  	return info.Value, nil
   136  }
   137  
   138  // SetCPUAffinity sets the processor affinity for the job object.
   139  // The affinity is passed in as a bitmask.
   140  func (job *JobObject) SetCPUAffinity(affinityBitMask uint64) error {
   141  	info, err := job.getExtendedInformation()
   142  	if err != nil {
   143  		return err
   144  	}
   145  	info.BasicLimitInformation.LimitFlags |= uint32(windows.JOB_OBJECT_LIMIT_AFFINITY)
   146  	info.BasicLimitInformation.Affinity = uintptr(affinityBitMask)
   147  	return job.setExtendedInformation(info)
   148  }
   149  
   150  // GetCPUAffinity gets the processor affinity for the job object.
   151  // The returned affinity is a bitmask.
   152  func (job *JobObject) GetCPUAffinity() (uint64, error) {
   153  	info, err := job.getExtendedInformation()
   154  	if err != nil {
   155  		return 0, err
   156  	}
   157  	return uint64(info.BasicLimitInformation.Affinity), nil
   158  }
   159  
   160  // SetIOLimit sets the IO limits specified on the job object.
   161  func (job *JobObject) SetIOLimit(maxBandwidth, maxIOPS int64) error {
   162  	ioInfo, err := job.getIOLimit()
   163  	if err != nil {
   164  		return err
   165  	}
   166  	ioInfo.ControlFlags |= winapi.JOB_OBJECT_IO_RATE_CONTROL_ENABLE
   167  	if maxBandwidth != 0 {
   168  		ioInfo.MaxBandwidth = maxBandwidth
   169  	}
   170  	if maxIOPS != 0 {
   171  		ioInfo.MaxIops = maxIOPS
   172  	}
   173  	return job.setIORateControlInfo(ioInfo)
   174  }
   175  
   176  // GetIOMaxBandwidthLimit gets the max bandwidth for the job object.
   177  func (job *JobObject) GetIOMaxBandwidthLimit() (int64, error) {
   178  	info, err := job.getIOLimit()
   179  	if err != nil {
   180  		return 0, err
   181  	}
   182  	return info.MaxBandwidth, nil
   183  }
   184  
   185  // GetIOMaxIopsLimit gets the max iops for the job object.
   186  func (job *JobObject) GetIOMaxIopsLimit() (int64, error) {
   187  	info, err := job.getIOLimit()
   188  	if err != nil {
   189  		return 0, err
   190  	}
   191  	return info.MaxIops, nil
   192  }
   193  
   194  // Helper function for getting a job object's extended information.
   195  func (job *JobObject) getExtendedInformation() (*windows.JOBOBJECT_EXTENDED_LIMIT_INFORMATION, error) {
   196  	job.handleLock.RLock()
   197  	defer job.handleLock.RUnlock()
   198  
   199  	if job.handle == 0 {
   200  		return nil, ErrAlreadyClosed
   201  	}
   202  
   203  	info := windows.JOBOBJECT_EXTENDED_LIMIT_INFORMATION{}
   204  	if err := winapi.QueryInformationJobObject(
   205  		job.handle,
   206  		windows.JobObjectExtendedLimitInformation,
   207  		unsafe.Pointer(&info),
   208  		uint32(unsafe.Sizeof(info)),
   209  		nil,
   210  	); err != nil {
   211  		return nil, fmt.Errorf("query %v returned error: %w", info, err)
   212  	}
   213  	return &info, nil
   214  }
   215  
   216  // Helper function for getting a job object's CPU rate control information.
   217  func (job *JobObject) getCPURateControlInformation() (*winapi.JOBOBJECT_CPU_RATE_CONTROL_INFORMATION, error) {
   218  	job.handleLock.RLock()
   219  	defer job.handleLock.RUnlock()
   220  
   221  	if job.handle == 0 {
   222  		return nil, ErrAlreadyClosed
   223  	}
   224  
   225  	info := winapi.JOBOBJECT_CPU_RATE_CONTROL_INFORMATION{}
   226  	if err := winapi.QueryInformationJobObject(
   227  		job.handle,
   228  		windows.JobObjectCpuRateControlInformation,
   229  		unsafe.Pointer(&info),
   230  		uint32(unsafe.Sizeof(info)),
   231  		nil,
   232  	); err != nil {
   233  		return nil, fmt.Errorf("query %v returned error: %w", info, err)
   234  	}
   235  	return &info, nil
   236  }
   237  
   238  // Helper function for setting a job object's extended information.
   239  func (job *JobObject) setExtendedInformation(info *windows.JOBOBJECT_EXTENDED_LIMIT_INFORMATION) error {
   240  	job.handleLock.RLock()
   241  	defer job.handleLock.RUnlock()
   242  
   243  	if job.handle == 0 {
   244  		return ErrAlreadyClosed
   245  	}
   246  
   247  	if _, err := windows.SetInformationJobObject(
   248  		job.handle,
   249  		windows.JobObjectExtendedLimitInformation,
   250  		uintptr(unsafe.Pointer(info)),
   251  		uint32(unsafe.Sizeof(*info)),
   252  	); err != nil {
   253  		return fmt.Errorf("failed to set Extended info %v on job object: %w", info, err)
   254  	}
   255  	return nil
   256  }
   257  
   258  // Helper function for querying job handle for IO limit information.
   259  func (job *JobObject) getIOLimit() (*winapi.JOBOBJECT_IO_RATE_CONTROL_INFORMATION, error) {
   260  	job.handleLock.RLock()
   261  	defer job.handleLock.RUnlock()
   262  
   263  	if job.handle == 0 {
   264  		return nil, ErrAlreadyClosed
   265  	}
   266  
   267  	ioInfo := &winapi.JOBOBJECT_IO_RATE_CONTROL_INFORMATION{}
   268  	var blockCount uint32 = 1
   269  
   270  	if _, err := winapi.QueryIoRateControlInformationJobObject(
   271  		job.handle,
   272  		nil,
   273  		&ioInfo,
   274  		&blockCount,
   275  	); err != nil {
   276  		return nil, fmt.Errorf("query %v returned error: %w", ioInfo, err)
   277  	}
   278  
   279  	if !isFlagSet(winapi.JOB_OBJECT_IO_RATE_CONTROL_ENABLE, ioInfo.ControlFlags) {
   280  		return nil, fmt.Errorf("query %v cannot get IO limits for job object without IO rate control option set", ioInfo)
   281  	}
   282  	return ioInfo, nil
   283  }
   284  
   285  // Helper function for setting a job object's IO rate control information.
   286  func (job *JobObject) setIORateControlInfo(ioInfo *winapi.JOBOBJECT_IO_RATE_CONTROL_INFORMATION) error {
   287  	job.handleLock.RLock()
   288  	defer job.handleLock.RUnlock()
   289  
   290  	if job.handle == 0 {
   291  		return ErrAlreadyClosed
   292  	}
   293  
   294  	if _, err := winapi.SetIoRateControlInformationJobObject(job.handle, ioInfo); err != nil {
   295  		return fmt.Errorf("failed to set IO limit info %v on job object: %w", ioInfo, err)
   296  	}
   297  	return nil
   298  }
   299  
   300  // Helper function for setting a job object's CPU rate control information.
   301  func (job *JobObject) setCPURateControlInfo(cpuInfo *winapi.JOBOBJECT_CPU_RATE_CONTROL_INFORMATION) error {
   302  	job.handleLock.RLock()
   303  	defer job.handleLock.RUnlock()
   304  
   305  	if job.handle == 0 {
   306  		return ErrAlreadyClosed
   307  	}
   308  	if _, err := windows.SetInformationJobObject(
   309  		job.handle,
   310  		windows.JobObjectCpuRateControlInformation,
   311  		uintptr(unsafe.Pointer(cpuInfo)),
   312  		uint32(unsafe.Sizeof(cpuInfo)),
   313  	); err != nil {
   314  		return fmt.Errorf("failed to set cpu limit info %v on job object: %w", cpuInfo, err)
   315  	}
   316  	return nil
   317  }
   318
View as plain text