github.com/castai/kvisor@v1.7.1-0.20240516114728-b3572a2607b5/cmd/controller/state/kubebench/delta.go (about)

     1  package kubebench
     2  
     3  import (
     4  	"sync"
     5  	"time"
     6  
     7  	"github.com/samber/lo"
     8  	corev1 "k8s.io/api/core/v1"
     9  	"k8s.io/apimachinery/pkg/types"
    10  	"k8s.io/apimachinery/pkg/util/wait"
    11  )
    12  
    13  func newDeltaState() *nodeDeltaState {
    14  	return &nodeDeltaState{
    15  		objectMap: make(map[types.UID]*nodeJob),
    16  		mu:        sync.Mutex{},
    17  	}
    18  }
    19  
    20  type nodeJob struct {
    21  	node *corev1.Node
    22  
    23  	backoff wait.Backoff
    24  	next    time.Time
    25  }
    26  
    27  func (n *nodeJob) ready() bool {
    28  	return n.next.Before(time.Now())
    29  }
    30  
    31  func (n *nodeJob) setFailed() {
    32  	n.next = time.Now().Add(n.backoff.Step())
    33  }
    34  
    35  type nodeDeltaState struct {
    36  	objectMap map[types.UID]*nodeJob
    37  	mu        sync.Mutex
    38  }
    39  
    40  func (d *nodeDeltaState) upsert(o *corev1.Node) {
    41  	d.mu.Lock()
    42  	defer d.mu.Unlock()
    43  
    44  	key := o.GetUID()
    45  	if job, ok := d.objectMap[key]; ok {
    46  		job.node = o
    47  		return
    48  	}
    49  
    50  	nodeBackoff := wait.Backoff{
    51  		Duration: time.Second * 15,
    52  		Factor:   3,
    53  		Steps:    8,
    54  	}
    55  	w := &nodeJob{
    56  		backoff: nodeBackoff,
    57  		node:    o,
    58  		next:    time.Now(),
    59  	}
    60  	d.objectMap[key] = w
    61  }
    62  
    63  func (d *nodeDeltaState) delete(o *corev1.Node) {
    64  	d.mu.Lock()
    65  	defer d.mu.Unlock()
    66  
    67  	delete(d.objectMap, o.GetUID())
    68  }
    69  
    70  func (d *nodeDeltaState) peek() []*nodeJob {
    71  	d.mu.Lock()
    72  	defer d.mu.Unlock()
    73  
    74  	return lo.Values(d.objectMap)
    75  }