github.com/kubeflow/training-operator@v1.7.0/pkg/util/counter.go (about)

     1  package util
     2  
     3  import (
     4  	"fmt"
     5  	"sync"
     6  )
     7  
     8  type Counter struct {
     9  	lock sync.Mutex
    10  	data map[string]int
    11  }
    12  
    13  func NewCounter() *Counter {
    14  	return &Counter{
    15  		lock: sync.Mutex{},
    16  		data: map[string]int{},
    17  	}
    18  }
    19  
    20  func (c *Counter) Inc(key string) {
    21  	c.lock.Lock()
    22  	defer c.lock.Unlock()
    23  
    24  	v, ok := c.data[key]
    25  	if ok {
    26  		c.data[key] = v + 1
    27  		return
    28  	}
    29  	c.data[key] = 0
    30  }
    31  
    32  func (c *Counter) DeleteKey(key string) {
    33  	c.lock.Lock()
    34  	defer c.lock.Lock()
    35  
    36  	delete(c.data, key)
    37  }
    38  
    39  func (c *Counter) Counts(key string) (int, error) {
    40  	c.lock.Lock()
    41  	defer c.lock.Unlock()
    42  
    43  	v, ok := c.data[key]
    44  	if !ok {
    45  		return 0, fmt.Errorf("cannot get key %s", key)
    46  	}
    47  	var err error = nil
    48  	if v < 0 {
    49  		err = fmt.Errorf("count %s:%d is negative", key, v)
    50  	}
    51  	return v, err
    52  }
    53  
    54  func (c *Counter) Dec(key string) error {
    55  	c.lock.Lock()
    56  	defer c.lock.Unlock()
    57  
    58  	v, ok := c.data[key]
    59  	if ok {
    60  		if v > 1 {
    61  			c.data[key] = v - 1
    62  			return nil
    63  		}
    64  		if v == 1 {
    65  			c.DeleteKey(key)
    66  			return nil
    67  		}
    68  		return fmt.Errorf("cannot minus one: key %s has value %d", key, v)
    69  	}
    70  	return fmt.Errorf("cannot find key %s", key)
    71  }