Fix unit test failure of semaphore and cluster

go.mod

@@ -46,6 +46,7 @@ require (
 	go.etcd.io/etcd/client/v3 v3.5.0
 	go.etcd.io/etcd/server/v3 v3.5.0
 	go.uber.org/zap v1.17.0
+	golang.org/x/sync v0.0.0-20210220032951-036812b2e83c
 	gopkg.in/yaml.v2 v2.4.0
 	k8s.io/api v0.20.7
 	k8s.io/apimachinery v0.20.7

pkg/cluster/cluster_test.go

@@ -33,11 +33,31 @@ func mockClusters(count int) []*cluster {
 		panic(fmt.Errorf("new cluster failed: %v", err))
 	}
 	clusters[0] = bootCluster.(*cluster)
+
 	time.Sleep(HeartbeatInterval)
 
 	for i := 1; i < count; i++ {
 		opts[i].ClusterJoinURLs = opts[0].ClusterListenPeerURLs
+
 		cls, err := New(opts[i])
+
+		if err != nil {
+			totalRetryTime := time.After(60 * time.Second)
+		Loop:
+			for {
+				if err == nil {
+					break
+				}
+				select {
+				case <-totalRetryTime:
+					break Loop
+
+				case <-time.After(HeartbeatInterval):
+					cls, err = New(opts[i])
+				}
+			}
+
+		}
 		if err != nil {
 			panic(fmt.Errorf("new cluster failed: %v", err))
 		}

pkg/cluster/member_test.go

@@ -19,6 +19,7 @@ package cluster
 
 import (
 	"fmt"
+	"io/ioutil"
 	"os"
 	"path/filepath"
 	"reflect"
@@ -33,13 +34,17 @@ import (
 	"github.com/megaease/easegress/pkg/option"
 )
 
-const tempDir = "/tmp/eg-test"
-
-var memberCounter = 0
+var (
+	memberCounter = 0
+	tempDir       = os.TempDir()
+)
 
 func TestMain(m *testing.M) {
-	absLogDir := filepath.Join(tempDir, "global-log")
-	os.MkdirAll(absLogDir, 0o755)
+	absLogDir, err := ioutil.TempDir(tempDir, "global-log")
+	if err != nil {
+		panic(fmt.Errorf("create tmp dir failed: %v", err))
+	}
+
 	logger.Init(&option.Options{
 		Name:      "member-for-log",
 		AbsLogDir: absLogDir,
@@ -49,7 +54,6 @@ func TestMain(m *testing.M) {
 
 	logger.Sync()
 	os.RemoveAll(tempDir)
-
 	os.Exit(code)
 }
 

pkg/object/httpserver/listen.go

@@ -72,7 +72,7 @@ func (l *LimitListener) Accept() (net.Conn, error) {
 
 // SetMaxConnection sets max connection.
 func (l *LimitListener) SetMaxConnection(n uint32) {
-	l.sem.SetMaxCount(n)
+	l.sem.SetMaxCount(int64(n))
 }
 
 // Close closes LimitListener.

pkg/util/sem/semaphore.go

@@ -17,69 +17,81 @@
 
 package sem
 
-import "sync"
+import (
+	"context"
+	"sync"
 
-const capacity = 2000000
+	"golang.org/x/sync/semaphore"
+)
 
+const maxCapacity int64 = 20_000_000
+
+//Semaphore supports to change the max sema amount at runtime.
+//  Semaphore employs golang.org/x/sync/semaphore.Weighted with a maxCapacity.
+//  And tuning the realCapacity by a Acquire and Release in the background.
+//  the realCapacity can not exceeds the maxCapacity.
 type Semaphore struct {
-	sem       uint32
-	lock      *sync.Mutex
-	guardChan chan *struct{}
+	sem          *semaphore.Weighted
+	lock         sync.Mutex
+	realCapacity int64
 }
 
 func NewSem(n uint32) *Semaphore {
 	s := &Semaphore{
-		sem:       n,
-		lock:      &sync.Mutex{},
-		guardChan: make(chan *struct{}, capacity),
+		sem:          semaphore.NewWeighted(maxCapacity),
+		realCapacity: int64(n),
 	}
 
-	go func() {
-		for i := uint32(0); i < n; i++ {
-			s.guardChan <- &struct{}{}
-		}
-	}()
-
+	s.sem.Acquire(context.Background(), maxCapacity-s.realCapacity)
 	return s
 }
 
 func (s *Semaphore) Acquire() {
-	<-s.guardChan
+	s.AcquireWithContext(context.Background())
+}
+
+func (s *Semaphore) AcquireWithContext(ctx context.Context) error {
+	return s.sem.Acquire(ctx, 1)
 }
 
 func (s *Semaphore) AcquireRaw() chan *struct{} {
-	return s.guardChan
+	s.AcquireWithContext(context.Background())
+	return make(chan *struct{})
 }
 
 func (s *Semaphore) Release() {
-	s.guardChan <- &struct{}{}
+	s.sem.Release(1)
 }
 
-func (s *Semaphore) SetMaxCount(n uint32) {
+func (s *Semaphore) SetMaxCount(n int64) (done chan struct{}) {
 	s.lock.Lock()
 	defer s.lock.Unlock()
 
-	if n > capacity {
-		n = capacity
+	if n > maxCapacity {
+		n = maxCapacity
 	}
 
-	if n == s.sem {
+	if n == s.realCapacity {
 		return
 	}
 
-	old := s.sem
-	s.sem = n
+	old := s.realCapacity
+	s.realCapacity = n
 
+	done = make(chan struct{})
 	go func() {
 		if n > old {
-			for i := uint32(0); i < n-old; i++ {
-				s.guardChan <- &struct{}{}
+			for i := int64(0); i < n-old; i++ {
+				s.sem.Release(1)
+			}
+		} else {
+			for i := int64(0); i < old-n; i++ {
+				s.sem.Acquire(context.Background(), 1)
 			}
-			return
 		}
 
-		for i := uint32(0); i < old-n; i++ {
-			<-s.guardChan
-		}
+		done <- struct{}{}
 	}()
+
+	return done
 }

pkg/util/sem/semaphore_test.go

@@ -18,16 +18,12 @@
 package sem
 
 import (
-	"math/rand"
+	"context"
 	"sync"
 	"testing"
 	"time"
 )
 
-type Case struct {
-	maxSem int
-}
-
 func TestSemaphore0(t *testing.T) {
 	s := NewSem(0)
 
@@ -44,112 +40,41 @@ func TestSemaphore0(t *testing.T) {
 	}
 }
 
-func TestSemaphoreRobust(t *testing.T) {
+func TestSemaphoreChangeAtRuntime(t *testing.T) {
 	s := NewSem(10)
-	w := &sync.WaitGroup{}
-	w.Add(101)
-	// change maxSem randomly
-	go func() {
-		for i := 0; i < 500; i++ {
-			time.Sleep(1 * time.Millisecond)
-			s.SetMaxCount(uint32(rand.Intn(100000)))
-		}
-		w.Done()
-	}()
-
-	for x := 0; x < 100; x++ {
-		go func() {
-			for j := 0; j < 100; j++ {
-				s.Acquire()
-				time.Sleep(5 * time.Millisecond)
-				s.Release()
-			}
-			w.Done()
-		}()
-	}
-	w.Wait()
-
-	// confirm it still works
-	s.SetMaxCount(20)
-	time.Sleep(10 * time.Millisecond)
-
-	runCase(s, &Case{23}, t)
-}
-
-func TestSemaphoreN(t *testing.T) {
-	s := NewSem(20)
-	Cases := []*Case{
-		{maxSem: 37},
-		{maxSem: 45},
-		{maxSem: 3},
-		{maxSem: 1000},
-		{maxSem: 235},
-		{maxSem: 800},
-		{maxSem: 587},
-	}
 
-	for _, c := range Cases {
+	testcases := []int64{40, 2}
+	for _, c := range testcases {
 		runCase(s, c, t)
-	}
-}
 
-func BenchmarkSemaphore(b *testing.B) {
-	s := NewSem(uint32(b.N/2 + 1))
-	for i := 0; i < b.N; i++ {
-		b.RunParallel(func(pb *testing.PB) {
-			for pb.Next() {
-				s.Acquire()
-				s.Release()
-			}
-		})
 	}
+
 }
 
-func runCase(s *Semaphore, c *Case, t *testing.T) {
-	s.SetMaxCount(uint32(c.maxSem))
-	time.Sleep(10 * time.Millisecond)
+func runCase(s *Semaphore, maxCount int64, t *testing.T) {
+	// step 0, change max count at runtime
+	done := s.SetMaxCount(maxCount)
+	<- done
 
-	w := &sync.WaitGroup{}
-	w.Add(c.maxSem)
-	for i := 0; i < c.maxSem; i++ {
-		go func(w *sync.WaitGroup) {
+	// step 1, try to acquire max count, should be OK
+	wg := &sync.WaitGroup{}
+	wg.Add(int(maxCount))
+	for i := 0; i < int(maxCount); i++ {
+		go func() {
 			s.Acquire()
-			defer s.Release()
-			time.Sleep(100 * time.Millisecond)
-			w.Done()
-		}(w)
-	}
-	begin := time.Now()
-	w.Wait()
-	d := time.Since(begin)
-	if d > 100*time.Millisecond+10*time.Millisecond {
-		t.Errorf("time too long: %v, sem: %d, trans: %d", d, c.maxSem, c.maxSem)
-	}
-	if d < 100*time.Millisecond-10*time.Millisecond {
-		t.Errorf("time too short: %v, sem: %d, trans: %d", d, c.maxSem, c.maxSem)
+			wg.Done()
+		}()
 	}
+	wg.Wait()
 
-	s.SetMaxCount(uint32(c.maxSem - 1))
-	time.Sleep(10 * time.Millisecond)
-
-	w = &sync.WaitGroup{}
-	w.Add(c.maxSem)
-	for i := 0; i < c.maxSem; i++ {
-		go func(w *sync.WaitGroup) {
-			s.Acquire()
-			defer s.Release()
-			time.Sleep(100 * time.Millisecond)
-			w.Done()
-		}(w)
-	}
-	begin = time.Now()
-	w.Wait()
-	d = time.Since(begin)
-	if d < 200*time.Millisecond-10*time.Millisecond {
-		t.Errorf("time too short: %v, sem: %d, trans: %d", d, c.maxSem-1, c.maxSem)
+	// step 2: try to acquire one more, should timeout
+	ctx, _ := context.WithTimeout(context.Background(), 100*time.Millisecond)
+	err := s.AcquireWithContext(ctx)
+	if err == nil {
+		t.Fatalf("sema count exceeds the maxCount: %d", maxCount)
 	}
 
-	if d > 200*time.Millisecond+10*time.Millisecond {
-		t.Errorf("time too long: %v, sem: %d, trans: %d", d, c.maxSem-1, c.maxSem)
+	for i := 0; i < int(maxCount); i++ {
+		s.Release()
 	}
 }