github.com/sdibtacm/sandbox@v0.0.0-20200320120712-60470cf803dc/exec/limiter.go (about)

     1  //+build linux
     2  
     3  package exec
     4  
     5  /*
     6  #include <unistd.h>
     7  int tickPreSec() {
     8  	return sysconf(_SC_CLK_TCK);
     9  }
    10  */
    11  import "C"
    12  
    13  import (
    14  	"github.com/sdibtacm/sandbox/exec/log"
    15  	"github.com/sdibtacm/sandbox/units/pstree"
    16  	"time"
    17  )
    18  
    19  var tickPerSec int
    20  var kernelTimeMod uint
    21  
    22  func init() {
    23  	tickPerSec = int(C.tickPreSec())
    24  	kernelTimeMod = uint(1000 / float64(tickPerSec))
    25  }
    26  
    27  func (c *Cmd) limiter() {
    28  	var used Resource
    29  	for !c.Process.Done() {
    30  		_ = calcUsed(&used, c.Process.Pid)
    31  		//time.Sleep(5 * time.Microsecond)
    32  		c.resourceStats.ClockTime = uint((time.Now().UnixNano() - c.startTimestamp.UnixNano()) / 1e6)
    33  		c.resourceStats.Thread = used.Thread
    34  		c.resourceStats.Memory = used.Memory
    35  		c.resourceStats.CpuTime = used.CpuTime
    36  		used.Memory = 0
    37  		used.Thread = 0
    38  		used.CpuTime = 0
    39  		go c.updateMaxResource()
    40  		go c.checkResource()
    41  	}
    42  }
    43  
    44  func (c *Cmd) checkResource() {
    45  	if c.ResourceLimit.Memory != BYTE_UNRESOURCE && c.resourceMaxStats.Memory > c.ResourceLimit.Memory {
    46  		_ = c.Process.KillGroup()
    47  	}
    48  	if c.resourceMaxStats.Thread > c.ResourceLimit.Thread {
    49  		_ = c.Process.KillGroup()
    50  	}
    51  	if c.ResourceLimit.CpuTime != TIME_UNRESOURCE && c.resourceMaxStats.CpuTime > c.ResourceLimit.CpuTime {
    52  		_ = c.Process.KillGroup()
    53  	}
    54  }
    55  
    56  func (c *Cmd) updateMaxResource() {
    57  	if c.resourceMaxStats.Memory < c.resourceStats.Memory {
    58  		c.resourceMaxStats.Memory = c.resourceStats.Memory
    59  	}
    60  	if c.resourceMaxStats.Thread < c.resourceStats.Thread {
    61  		c.resourceMaxStats.Thread = c.resourceStats.Thread
    62  	}
    63  	if c.resourceMaxStats.CpuTime < c.resourceStats.CpuTime {
    64  		c.resourceMaxStats.CpuTime = c.resourceStats.CpuTime
    65  	}
    66  }
    67  
    68  func calcUsed(r *Resource, rootPid int) error {
    69  	pt, err := pstree.New()
    70  	if err != nil {
    71  		log.GetLog().Warning("pstree scan error: {}", err)
    72  		return err
    73  	}
    74  	procs := pt.Procs
    75  	pids := []int{rootPid}
    76  	//r.CpuTime = uint(procs[rootPid].Stat.Stime + procs[rootPid].Stat.Utime + procs[rootPid].Stat.Cstime + procs[rootPid].Stat.Cutime) * kernelTimeMod
    77  	for i := 0; i < len(pids); i++ {
    78  		pid := pids[i]
    79  		if len(procs[pid].Children) > 0 {
    80  			pids = append(pids, procs[pid].Children...)
    81  		}
    82  		r.CpuTime += uint(procs[pid].Stat.Stime+procs[pid].Stat.Utime) * kernelTimeMod
    83  		r.Memory += procs[pid].Stat.Vsize
    84  		r.Thread += uint(procs[pid].Stat.Nthreads)
    85  	}
    86  	return nil
    87  }