[ws-daemon] Fix CPU limit annotation

components/ws-daemon/pkg/cpulimit/cpulimit.go

@@ -18,6 +18,9 @@ type Workspace struct {
 	NrThrottled uint64
 	Usage       CPUTime
 	QoS         int
+
+	BaseLimit  Bandwidth
+	BurstLimit Bandwidth
 }
 
 type WorkspaceHistory struct {
@@ -113,7 +116,7 @@ func (d *Distributor) Tick(dt time.Duration) (DistributorDebug, error) {
 		return DistributorDebug{}, err
 	}
 
-	f := make(map[string]struct{}, len(ws))
+	wsidx := make(map[string]Workspace, len(ws))
 	for _, w := range ws {
 		h, ok := d.History[w.ID]
 		if !ok {
@@ -123,10 +126,10 @@ func (d *Distributor) Tick(dt time.Duration) (DistributorDebug, error) {
 			d.History[w.ID] = h
 		}
 		h.Update(w)
-		f[w.ID] = struct{}{}
+		wsidx[w.ID] = w
 	}
 	for oldWS := range d.History {
-		if _, found := f[oldWS]; !found {
+		if _, found := wsidx[oldWS]; !found {
 			delete(d.History, oldWS)
 		}
 	}
@@ -172,15 +175,23 @@ func (d *Distributor) Tick(dt time.Duration) (DistributorDebug, error) {
 	var burstBandwidth Bandwidth
 	for _, id := range wsOrder {
 		ws := d.History[id]
-		limit := d.Limiter.Limit(ws.Usage())
+		limiter := d.Limiter
+		if w := wsidx[id]; w.BaseLimit > 0 {
+			limiter = FixedLimiter(w.BaseLimit)
+		}
+		limit := limiter.Limit(ws.Usage())
 
 		// if we didn't get the max bandwidth, but were throttled last time
 		// and there's still some bandwidth left to give, let's act as if had
 		// never spent any CPU time and assume the workspace will spend their
 		// entire bandwidth at once.
 		var burst bool
 		if totalBandwidth < d.TotalBandwidth && ws.Throttled() {
-			limit = d.BurstLimiter.Limit(ws.Usage())
+			limiter := d.BurstLimiter
+			if w := wsidx[id]; w.BaseLimit > 0 {
+				limiter = FixedLimiter(w.BurstLimit)
+			}
+			limit = limiter.Limit(ws.Usage())
 
 			// We assume the workspace is going to use as much as their limit allows.
 			// This might not be true, because their process which consumed so much CPU

components/ws-daemon/pkg/cpulimit/dispatch.go

@@ -98,9 +98,11 @@ type DispatchListener struct {
 }
 
 type workspace struct {
-	CFS       CFSController
-	OWI       logrus.Fields
-	HardLimit ResourceLimiter
+	CFS        CgroupCFSController
+	OWI        logrus.Fields
+	BaseLimit  Bandwidth
+	BurstLimit Bandwidth
+	HardLimit  ResourceLimiter
 
 	lastThrottled uint64
 }
@@ -135,6 +137,8 @@ func (d *DispatchListener) source(context.Context) ([]Workspace, error) {
 			ID:          id,
 			NrThrottled: throttled,
 			Usage:       usage,
+			BaseLimit:   w.BaseLimit,
+			BurstLimit:  w.BurstLimit,
 		})
 	}
 	return res, nil
@@ -215,7 +219,8 @@ func (d *DispatchListener) WorkspaceUpdated(ctx context.Context, ws *dispatch.Wo
 		if err != nil {
 			return xerrors.Errorf("cannot enforce fixed CPU limit: %w", err)
 		}
-		wsinfo.HardLimit = FixedLimiter(BandwidthFromQuantity(limit))
+		wsinfo.BaseLimit = BandwidthFromQuantity(limit)
+		wsinfo.BurstLimit = BandwidthFromQuantity(limit)
 	}
 
 	return nil