github.com/moleculer-go/moleculer@v0.3.3/test/counter.go (about)

     1  package test
     2  
     3  import (
     4  	"errors"
     5  	"fmt"
     6  	"sync"
     7  	"time"
     8  )
     9  
    10  var CounterCheckTimeout = 20 * time.Second
    11  
    12  func Counter() CounterCheck {
    13  	return CounterCheck{&sync.Mutex{}, make(map[string]int), make(map[string]int)}
    14  }
    15  
    16  type CounterCheck struct {
    17  	mutex             *sync.Mutex
    18  	observe           map[string]int
    19  	observeWithPrefix map[string]int
    20  }
    21  
    22  func (counter *CounterCheck) Inc(nodeID string, name string) {
    23  	go func() {
    24  		counter.mutex.Lock()
    25  		if value, exists := counter.observe[name]; exists {
    26  			counter.observe[name] = value + 1
    27  		} else {
    28  			counter.observe[name] = 1
    29  		}
    30  
    31  		prefixed := fmt.Sprint(name, "-", nodeID)
    32  		if value, exists := counter.observeWithPrefix[prefixed]; exists {
    33  			counter.observeWithPrefix[prefixed] = value + 1
    34  		} else {
    35  			counter.observeWithPrefix[prefixed] = 1
    36  		}
    37  		counter.mutex.Unlock()
    38  	}()
    39  }
    40  
    41  func (counter *CounterCheck) Clear() {
    42  	counter.mutex.Lock()
    43  	counter.observe = make(map[string]int)
    44  	counter.observeWithPrefix = make(map[string]int)
    45  	counter.mutex.Unlock()
    46  }
    47  
    48  func (counter *CounterCheck) CheckPrefixed(name string, value int) error {
    49  	return counter.checkAbs(&counter.observeWithPrefix, name, value)
    50  }
    51  
    52  func (counter *CounterCheck) Check(name string, value int) error {
    53  	return counter.checkAbs(&counter.observe, name, value)
    54  }
    55  
    56  func (counter *CounterCheck) checkAbs(values *map[string]int, name string, target int) error {
    57  	result := make(chan error)
    58  	go func() {
    59  		start := time.Now()
    60  		for {
    61  			counter.mutex.Lock()
    62  			value, exists := (*values)[name]
    63  			counter.mutex.Unlock()
    64  			if exists && value >= target || target == 0 {
    65  				result <- nil
    66  				break
    67  			}
    68  			if time.Since(start) > CounterCheckTimeout {
    69  				result <- errors.New(fmt.Sprint("counter check timed out! -> name: ", name, " target: ", target, " current: ", value))
    70  				break
    71  			}
    72  			time.Sleep(time.Microsecond)
    73  		}
    74  	}()
    75  	return <-result
    76  }