func()

in internal/resource/resource_windows.go [94:177]
63 lines of code
14 McCabe index (conditional complexity)

func (c *windowsClient) Apply(constraint Constraint) error {
	// If no constraints are set, return early.
	if constraint.MaxMemoryUsage == 0 && constraint.MaxCPUUsage == 0 {
		return nil
	}

	// Create a JobObject for this plugin.
	jobName, err := windows.UTF16PtrFromString(constraint.Name)
	if err != nil {
		return fmt.Errorf("failed to parse plugin name: %w", err)
	}
	jobObjectHandle, err := c.createJobObject(nil, jobName)
	if err != nil && err != windows.ERROR_ALREADY_EXISTS {
		return fmt.Errorf("error creating job object for plugin %s: %w", constraint.Name, err)
	}
	c.handles[constraint.Name] = jobObjectHandle
	galog.Debugf("Created JobObject for plugin %s with handle %v", constraint.Name, jobObjectHandle)

	// Get the process handle.
	procHandle, err := c.getProcessHandle(processAllAccess, false, uint32(constraint.PID))
	if err != nil {
		return fmt.Errorf("failed to open process handler for plugin %s: %w", constraint.Name, err)
	}
	defer windows.CloseHandle(procHandle)

	// Assign the process to the JobObject.
	if err := c.assignProcessToJobObject(jobObjectHandle, procHandle); err != nil {
		return fmt.Errorf("failed to assign process %v to job object: %w", constraint.Name, err)
	}

	// Set the process memory limits.
	if constraint.MaxMemoryUsage != 0 {
		// First validate that the MaxMemoryUsage is greater than the minimum
		// required, which is 20 pages.
		//
		// https://learn.microsoft.com/en-us/windows/win32/api/memoryapi/nf-memoryapi-setprocessworkingsetsizeex
		minimumSize := 20 * os.Getpagesize()
		if constraint.MaxMemoryUsage < int64(minimumSize) {
			return fmt.Errorf("MaxMemoryUsage %d is less than the minimum required memory %d", constraint.MaxMemoryUsage, minimumSize)
		}

		jobMemoryLimitInfoClass := uint32(windows.JobObjectBasicLimitInformation)
		jobLimitInfo := windows.JOBOBJECT_BASIC_LIMIT_INFORMATION{
			LimitFlags:            windows.JOB_OBJECT_LIMIT_WORKINGSET,
			MinimumWorkingSetSize: uintptr(minimumSize),
			MaximumWorkingSetSize: uintptr(constraint.MaxMemoryUsage),
		}
		if _, err = c.setInformationJobObject(
			jobObjectHandle,
			jobMemoryLimitInfoClass,
			uintptr(unsafe.Pointer(&jobLimitInfo)),
			uint32(unsafe.Sizeof(jobLimitInfo)),
		); err != nil {
			return fmt.Errorf("failed to set job object memory limits for plugin %s: %w", constraint.Name, err)
		}
	}

	if constraint.MaxCPUUsage != 0 {
		// Make sure the CPU usage set is realistic.
		if constraint.MaxCPUUsage > 100 {
			return fmt.Errorf("MaxCPUUsage %d exceeds 100%%", constraint.MaxCPUUsage)
		}

		// Set CPU rate limits. Usage limits is the maximum number of CPU cycles
		// allowed per 10,000 cycles.
		cpuCycles := uint32(constraint.MaxCPUUsage * 100)
		jobCPURateLimitInfoClass := uint32(windows.JobObjectCpuRateControlInformation)
		jobCPURateControl := jobObjectCPURateControlInfo{
			ControlFlag: jobObjectCPUHardCap | jobObjectCPURateControlEnable,
			Value:       cpuCycles,
		}

		if _, err = c.setInformationJobObject(
			jobObjectHandle,
			jobCPURateLimitInfoClass,
			uintptr(unsafe.Pointer(&jobCPURateControl)),
			uint32(unsafe.Sizeof(jobCPURateControl)),
		); err != nil {
			return fmt.Errorf("failed to set CPU rate limits for plugin %s: %w", constraint.Name, err)
		}
	}

	return nil
}