github.com/fufuok/balancer@v1.0.0/swrr_test.go (about)

     1  package balancer
     2  
     3  import (
     4  	"sync"
     5  	"sync/atomic"
     6  	"testing"
     7  )
     8  
     9  func TestSmoothWeightedRoundRobin(t *testing.T) {
    10  	lb := NewSmoothWeightedRoundRobin()
    11  	item := lb.Select()
    12  	if item != "" {
    13  		t.Fatalf("swrr expected empty, actual %s", item)
    14  	}
    15  
    16  	lb.Add("A", 0)
    17  	item = lb.Select()
    18  	if item != "" {
    19  		t.Fatalf("swrr expected empty, actual %s", item)
    20  	}
    21  	lb.Add("B", 1)
    22  	item = lb.Select()
    23  	if item != "B" {
    24  		t.Fatalf("swrr expected B, actual %s", item)
    25  	}
    26  
    27  	nodes := map[string]int{
    28  		"A": 0,
    29  		"B": 1,
    30  		"C": 7,
    31  		"D": 2,
    32  	}
    33  	lb = NewSmoothWeightedRoundRobin(nodes)
    34  	count := make(map[string]int)
    35  	for i := 0; i < 1000; i++ {
    36  		item := lb.Select()
    37  		count[item]++
    38  	}
    39  	if count["A"] != 0 || count["B"] != 100 || count["C"] != 700 || count["D"] != 200 {
    40  		t.Fatal("swrr wrong")
    41  	}
    42  
    43  	lb.Reset()
    44  
    45  	lb.Add("E", 10)
    46  	for i := 0; i < 2000; i++ {
    47  		item := lb.Select()
    48  		count[item]++
    49  	}
    50  	if count["A"] != 0 || count["B"] != 200 || count["C"] != 1400 || count["D"] != 400 || count["E"] != 1000 {
    51  		t.Fatal("swrr reset() wrong")
    52  	}
    53  
    54  	ok := lb.Remove("E")
    55  	if ok != true {
    56  		t.Fatal("swrr remove() wrong")
    57  	}
    58  
    59  	lb.Reset()
    60  
    61  	for i := 0; i < 1000; i++ {
    62  		item := lb.Select()
    63  		count[item]++
    64  	}
    65  	if count["A"] != 0 || count["B"] != 300 || count["C"] != 2100 || count["D"] != 600 {
    66  		t.Fatal("swrr wrong")
    67  	}
    68  
    69  	lb.RemoveAll()
    70  	lb.Add("F", 2)
    71  	lb.Add("F", 1)
    72  	all, ok := lb.All().(map[string]int)
    73  	if !ok || all["F"] != 1 {
    74  		t.Fatal("swrr all() wrong")
    75  	}
    76  
    77  	nodes = map[string]int{
    78  		"X": 0,
    79  		"Y": 1,
    80  	}
    81  	ok = lb.Update(nodes)
    82  	if ok != true {
    83  		t.Fatal("swrr update wrong")
    84  	}
    85  	all, ok = lb.All().(map[string]int)
    86  	if !ok || all["Y"] != 1 {
    87  		t.Fatal("swrr all() wrong")
    88  	}
    89  	item = lb.Select()
    90  	if item != "Y" {
    91  		t.Fatal("swrr update wrong")
    92  	}
    93  }
    94  
    95  func TestSmoothWeightedRoundRobin_C(t *testing.T) {
    96  	var (
    97  		a, b, c, d int64
    98  	)
    99  	nodes := map[string]int{
   100  		"A": 5,
   101  		"B": 1,
   102  		"C": 4,
   103  		"D": 0,
   104  	}
   105  	lb := NewSmoothWeightedRoundRobin(nodes)
   106  
   107  	var wg sync.WaitGroup
   108  	for i := 0; i < 500; i++ {
   109  		wg.Add(1)
   110  		go func() {
   111  			defer wg.Done()
   112  			for j := 0; j < 2000; j++ {
   113  				switch lb.Select() {
   114  				case "A":
   115  					atomic.AddInt64(&a, 1)
   116  				case "B":
   117  					atomic.AddInt64(&b, 1)
   118  				case "C":
   119  					atomic.AddInt64(&c, 1)
   120  				case "D":
   121  					atomic.AddInt64(&d, 1)
   122  				}
   123  			}
   124  		}()
   125  	}
   126  	wg.Wait()
   127  
   128  	if atomic.LoadInt64(&a) != 500000 {
   129  		t.Fatal("swrr wrong: a")
   130  	}
   131  	if atomic.LoadInt64(&b) != 100000 {
   132  		t.Fatal("swrr wrong: b")
   133  	}
   134  	if atomic.LoadInt64(&c) != 400000 {
   135  		t.Fatal("swrr wrong: c")
   136  	}
   137  	if atomic.LoadInt64(&d) != 0 {
   138  		t.Fatal("swrr wrong: d")
   139  	}
   140  }