@@ -25,18 +25,30 @@ package resourcetuner
2525import (
2626 "context"
2727 "errors"
28+ "runtime"
2829 "sync"
2930 "time"
3031
3132 "github.com/shirou/gopsutil/v4/cpu"
3233 "github.com/shirou/gopsutil/v4/mem"
3334 "go.einride.tech/pid"
35+ "go.temporal.io/sdk/log"
3436 "go.temporal.io/sdk/worker"
3537)
3638
3739type ResourceBasedTunerOptions struct {
40+ // TargetMem is the target overall system memory usage as value 0 and 1 that the controller will
41+ // attempt to maintain. Must be set nonzero.
3842 TargetMem float64
43+ // TargetCpu is the target overall system CPU usage as value 0 and 1 that the controller will
44+ // attempt to maintain. Must be set nonzero.
3945 TargetCpu float64
46+ // Passed to ResourceBasedSlotSupplierOptions.RampThrottle for activities.
47+ // If not set, the default value is 50ms.
48+ ActivityRampThrottle time.Duration
49+ // Passed to ResourceBasedSlotSupplierOptions.RampThrottle for workflows.
50+ // If not set, the default value is 0ms.
51+ WorkflowRampThrottle time.Duration
4052}
4153
4254// NewResourceBasedTuner creates a WorkerTuner that dynamically adjusts the number of slots based
@@ -50,10 +62,19 @@ func NewResourceBasedTuner(opts ResourceBasedTunerOptions) (worker.WorkerTuner,
5062 controller := NewResourceController (options )
5163 wfSS := & ResourceBasedSlotSupplier {controller : controller ,
5264 options : defaultWorkflowResourceBasedSlotSupplierOptions ()}
65+ if opts .WorkflowRampThrottle != 0 {
66+ wfSS .options .RampThrottle = opts .WorkflowRampThrottle
67+ }
5368 actSS := & ResourceBasedSlotSupplier {controller : controller ,
5469 options : defaultActivityResourceBasedSlotSupplierOptions ()}
70+ if opts .ActivityRampThrottle != 0 {
71+ actSS .options .RampThrottle = opts .ActivityRampThrottle
72+ }
5573 laSS := & ResourceBasedSlotSupplier {controller : controller ,
5674 options : defaultActivityResourceBasedSlotSupplierOptions ()}
75+ if opts .ActivityRampThrottle != 0 {
76+ laSS .options .RampThrottle = opts .ActivityRampThrottle
77+ }
5778 nexusSS := & ResourceBasedSlotSupplier {controller : controller ,
5879 options : defaultWorkflowResourceBasedSlotSupplierOptions ()}
5980 compositeTuner , err := worker .NewCompositeTuner (worker.CompositeTunerOptions {
@@ -163,7 +184,7 @@ func (r *ResourceBasedSlotSupplier) TryReserveSlot(info worker.SlotReservationIn
163184 numIssued := info .NumIssuedSlots ()
164185 if numIssued < r .options .MinSlots || (numIssued < r .options .MaxSlots &&
165186 time .Since (r .lastSlotIssuedAt ) > r .options .RampThrottle ) {
166- decision , err := r .controller .pidDecision ()
187+ decision , err := r .controller .pidDecision (info . Logger () )
167188 if err != nil {
168189 info .Logger ().Error ("Error calculating resource usage" , "error" , err )
169190 return nil
@@ -188,10 +209,14 @@ func (r *ResourceBasedSlotSupplier) MaxSlots() int {
188209type SystemInfoSupplier interface {
189210 // GetMemoryUsage returns the current system memory usage as a fraction of total memory between
190211 // 0 and 1.
191- GetMemoryUsage () (float64 , error )
212+ GetMemoryUsage (infoContext * SystemInfoContext ) (float64 , error )
192213 // GetCpuUsage returns the current system CPU usage as a fraction of total CPU usage between 0
193214 // and 1.
194- GetCpuUsage () (float64 , error )
215+ GetCpuUsage (infoContext * SystemInfoContext ) (float64 , error )
216+ }
217+
218+ type SystemInfoContext struct {
219+ Logger log.Logger
195220}
196221
197222// ResourceControllerOptions contains configurable parameters for a ResourceController.
@@ -262,7 +287,9 @@ type ResourceController struct {
262287func NewResourceController (options ResourceControllerOptions ) * ResourceController {
263288 var infoSupplier SystemInfoSupplier
264289 if options .InfoSupplier == nil {
265- infoSupplier = & psUtilSystemInfoSupplier {}
290+ infoSupplier = & psUtilSystemInfoSupplier {
291+ cGroupInfo : newCGroupInfo (),
292+ }
266293 } else {
267294 infoSupplier = options .InfoSupplier
268295 }
@@ -286,23 +313,22 @@ func NewResourceController(options ResourceControllerOptions) *ResourceControlle
286313 }
287314}
288315
289- func (rc * ResourceController ) pidDecision () (bool , error ) {
316+ func (rc * ResourceController ) pidDecision (logger log. Logger ) (bool , error ) {
290317 rc .mu .Lock ()
291318 defer rc .mu .Unlock ()
292319
293- memUsage , err := rc .infoSupplier .GetMemoryUsage ()
320+ memUsage , err := rc .infoSupplier .GetMemoryUsage (& SystemInfoContext { Logger : logger } )
294321 if err != nil {
295322 return false , err
296323 }
297- cpuUsage , err := rc .infoSupplier .GetCpuUsage ()
324+ cpuUsage , err := rc .infoSupplier .GetCpuUsage (& SystemInfoContext { Logger : logger } )
298325 if err != nil {
299326 return false , err
300327 }
301328 if memUsage >= rc .options .MemTargetPercent {
302329 // Never allow going over the memory target
303330 return false , nil
304331 }
305- //fmt.Printf("mem: %f, cpu: %f\n", memUsage, cpuUsage)
306332 elapsedTime := time .Since (rc .lastRefresh )
307333 // This shouldn't be possible with real implementations, but if the elapsed time is 0 the
308334 // PID controller can produce NaNs.
@@ -326,27 +352,54 @@ func (rc *ResourceController) pidDecision() (bool, error) {
326352}
327353
328354type psUtilSystemInfoSupplier struct {
329- mu sync.Mutex
355+ logger log.Logger
356+ mu sync.Mutex
357+ lastRefresh time.Time
358+
330359 lastMemStat * mem.VirtualMemoryStat
331360 lastCpuUsage float64
332- lastRefresh time.Time
361+
362+ stopTryingToGetCGroupInfo bool
363+ cGroupInfo cGroupInfo
333364}
334365
335- func (p * psUtilSystemInfoSupplier ) GetMemoryUsage () (float64 , error ) {
336- if err := p .maybeRefresh (); err != nil {
366+ type cGroupInfo interface {
367+ // Update requests an update of the cgroup stats. This is a no-op if not in a cgroup. Returns
368+ // true if cgroup stats should continue to be updated, false if not in a cgroup or the returned
369+ // error is considered unrecoverable.
370+ Update () (bool , error )
371+ // GetLastMemUsage returns last known memory usage as a fraction of the cgroup limit. 0 if not
372+ // in a cgroup or limit is not set.
373+ GetLastMemUsage () float64
374+ // GetLastCPUUsage returns last known CPU usage as a fraction of the cgroup limit. 0 if not in a
375+ // cgroup or limit is not set.
376+ GetLastCPUUsage () float64
377+ }
378+
379+ func (p * psUtilSystemInfoSupplier ) GetMemoryUsage (infoContext * SystemInfoContext ) (float64 , error ) {
380+ if err := p .maybeRefresh (infoContext ); err != nil {
337381 return 0 , err
338382 }
383+ lastCGroupMem := p .cGroupInfo .GetLastMemUsage ()
384+ if lastCGroupMem != 0 {
385+ return lastCGroupMem , nil
386+ }
339387 return p .lastMemStat .UsedPercent / 100 , nil
340388}
341389
342- func (p * psUtilSystemInfoSupplier ) GetCpuUsage () (float64 , error ) {
343- if err := p .maybeRefresh (); err != nil {
390+ func (p * psUtilSystemInfoSupplier ) GetCpuUsage (infoContext * SystemInfoContext ) (float64 , error ) {
391+ if err := p .maybeRefresh (infoContext ); err != nil {
344392 return 0 , err
345393 }
394+
395+ lastCGroupCPU := p .cGroupInfo .GetLastCPUUsage ()
396+ if lastCGroupCPU != 0 {
397+ return lastCGroupCPU , nil
398+ }
346399 return p .lastCpuUsage / 100 , nil
347400}
348401
349- func (p * psUtilSystemInfoSupplier ) maybeRefresh () error {
402+ func (p * psUtilSystemInfoSupplier ) maybeRefresh (infoContext * SystemInfoContext ) error {
350403 if time .Since (p .lastRefresh ) < 100 * time .Millisecond {
351404 return nil
352405 }
@@ -360,16 +413,24 @@ func (p *psUtilSystemInfoSupplier) maybeRefresh() error {
360413 defer cancelFn ()
361414 memStat , err := mem .VirtualMemoryWithContext (ctx )
362415 if err != nil {
363- println ("Refresh error: " , err )
364416 return err
365417 }
366418 cpuUsage , err := cpu .PercentWithContext (ctx , 0 , false )
367419 if err != nil {
368- println ("Refresh error: " , err )
369420 return err
370421 }
422+
371423 p .lastMemStat = memStat
372424 p .lastCpuUsage = cpuUsage [0 ]
425+
426+ if runtime .GOOS == "linux" && ! p .stopTryingToGetCGroupInfo {
427+ continueUpdates , err := p .cGroupInfo .Update ()
428+ if err != nil {
429+ infoContext .Logger .Warn ("Failed to get cgroup stats" , "error" , err )
430+ }
431+ p .stopTryingToGetCGroupInfo = ! continueUpdates
432+ }
433+
373434 p .lastRefresh = time .Now ()
374435 return nil
375436}
0 commit comments