Refactor port allocation logic a little, deflake tests.

pkg/proxy/userspace/port_allocator.go

@@ -57,7 +57,7 @@ func newPortAllocator(r net.PortRange) PortAllocator {
 	if r.Base == 0 {
 		return &randomAllocator{}
 	}
-	return newPortRangeAllocator(r)
+	return newPortRangeAllocator(r, true)
 }
 
 const (
@@ -74,7 +74,7 @@ type rangeAllocator struct {
 	rand  *rand.Rand
 }
 
-func newPortRangeAllocator(r net.PortRange) PortAllocator {
+func newPortRangeAllocator(r net.PortRange, autoFill bool) PortAllocator {
 	if r.Base == 0 || r.Size == 0 {
 		panic("illegal argument: may not specify an empty port range")
 	}
@@ -83,24 +83,29 @@ func newPortRangeAllocator(r net.PortRange) PortAllocator {
 		ports:     make(chan int, portsBufSize),
 		rand:      rand.New(rand.NewSource(time.Now().UnixNano())),
 	}
-	go wait.Until(func() { ra.fillPorts(wait.NeverStop) }, nextFreePortCooldown, wait.NeverStop)
+	if autoFill {
+		go wait.Forever(func() { ra.fillPorts() }, nextFreePortCooldown)
+	}
 	return ra
 }
 
 // fillPorts loops, always searching for the next free port and, if found, fills the ports buffer with it.
-// this func blocks until either there are no remaining free ports, or else the stopCh chan is closed.
-func (r *rangeAllocator) fillPorts(stopCh <-chan struct{}) {
+// this func blocks unless there are no remaining free ports.
+func (r *rangeAllocator) fillPorts() {
 	for {
-		port := r.nextFreePort()
-		if port == -1 {
+		if !r.fillPortsOnce() {
 			return
 		}
-		select {
-		case <-stopCh:
-			return
-		case r.ports <- port:
 	}
+}
+
+func (r *rangeAllocator) fillPortsOnce() bool {
+	port := r.nextFreePort()
+	if port == -1 {
+		return false
 	}
+	r.ports <- port
+	return true
 }
 
 // nextFreePort finds a free port, first picking a random port. if that port is already in use

pkg/proxy/userspace/port_allocator_test.go

@@ -31,15 +31,26 @@ func TestRangeAllocatorEmpty(t *testing.T) {
 			t.Fatalf("expected panic because of empty port range: %#v", r)
 		}
 	}()
-	_ = newPortRangeAllocator(*r)
+	_ = newPortRangeAllocator(*r, true)
 }
 
 func TestRangeAllocatorFullyAllocated(t *testing.T) {
 	r := &net.PortRange{}
 	r.Set("1-1")
-	pra := newPortRangeAllocator(*r)
+	// Don't auto-fill ports, we'll manually turn the crank
+	pra := newPortRangeAllocator(*r, false)
 	a := pra.(*rangeAllocator)
 
+	// Fill in the one available port
+	if !a.fillPortsOnce() {
+		t.Fatalf("Expected to be able to fill ports")
+	}
+
+	// There should be no ports available
+	if a.fillPortsOnce() {
+		t.Fatalf("Expected to be unable to fill ports")
+	}
+
 	p, err := a.AllocateNext()
 	if err != nil {
 		t.Fatalf("unexpected error: %v", err)
@@ -68,6 +79,11 @@ func TestRangeAllocatorFullyAllocated(t *testing.T) {
 	}
 	a.lock.Unlock()
 
+	// Fill in the one available port
+	if !a.fillPortsOnce() {
+		t.Fatalf("Expected to be able to fill ports")
+	}
+
 	p, err = a.AllocateNext()
 	if err != nil {
 		t.Fatalf("unexpected error: %v", err)
@@ -91,13 +107,16 @@ func TestRangeAllocatorFullyAllocated(t *testing.T) {
 func TestRangeAllocator_RandomishAllocation(t *testing.T) {
 	r := &net.PortRange{}
 	r.Set("1-100")
-	pra := newPortRangeAllocator(*r)
+	pra := newPortRangeAllocator(*r, false)
 	a := pra.(*rangeAllocator)
 
 	// allocate all the ports
 	var err error
 	ports := make([]int, 100, 100)
 	for i := 0; i < 100; i++ {
+		if !a.fillPortsOnce() {
+			t.Fatalf("Expected to be able to fill ports")
+		}
 		ports[i], err = a.AllocateNext()
 		if err != nil {
 			t.Fatalf("unexpected error: %v", err)
@@ -113,6 +132,10 @@ func TestRangeAllocator_RandomishAllocation(t *testing.T) {
 		a.lock.Unlock()
 	}
 
+	if a.fillPortsOnce() {
+		t.Fatalf("Expected to be unable to fill ports")
+	}
+
 	// release them all
 	for i := 0; i < 100; i++ {
 		a.Release(ports[i])
@@ -127,6 +150,9 @@ func TestRangeAllocator_RandomishAllocation(t *testing.T) {
 	// allocate the ports again
 	rports := make([]int, 100, 100)
 	for i := 0; i < 100; i++ {
+		if !a.fillPortsOnce() {
+			t.Fatalf("Expected to be able to fill ports")
+		}
 		rports[i], err = a.AllocateNext()
 		if err != nil {
 			t.Fatalf("unexpected error: %v", err)
@@ -142,6 +168,10 @@ func TestRangeAllocator_RandomishAllocation(t *testing.T) {
 		a.lock.Unlock()
 	}
 
+	if a.fillPortsOnce() {
+		t.Fatalf("Expected to be unable to fill ports")
+	}
+
 	if reflect.DeepEqual(ports, rports) {
 		t.Fatalf("expected re-allocated ports to be in a somewhat random order")
 	}

pkg/proxy/winuserspace/port_allocator.go

@@ -57,7 +57,7 @@ func newPortAllocator(r net.PortRange) PortAllocator {
 	if r.Base == 0 {
 		return &randomAllocator{}
 	}
-	return newPortRangeAllocator(r)
+	return newPortRangeAllocator(r, true)
 }
 
 const (
@@ -74,7 +74,7 @@ type rangeAllocator struct {
 	rand  *rand.Rand
 }
 
-func newPortRangeAllocator(r net.PortRange) PortAllocator {
+func newPortRangeAllocator(r net.PortRange, autoFill bool) PortAllocator {
 	if r.Base == 0 || r.Size == 0 {
 		panic("illegal argument: may not specify an empty port range")
 	}
@@ -83,24 +83,29 @@ func newPortRangeAllocator(r net.PortRange) PortAllocator {
 		ports:     make(chan int, portsBufSize),
 		rand:      rand.New(rand.NewSource(time.Now().UnixNano())),
 	}
-	go wait.Until(func() { ra.fillPorts(wait.NeverStop) }, nextFreePortCooldown, wait.NeverStop)
+	if autoFill {
+		go wait.Forever(func() { ra.fillPorts() }, nextFreePortCooldown)
+	}
 	return ra
 }
 
 // fillPorts loops, always searching for the next free port and, if found, fills the ports buffer with it.
-// this func blocks until either there are no remaining free ports, or else the stopCh chan is closed.
-func (r *rangeAllocator) fillPorts(stopCh <-chan struct{}) {
+// this func blocks unless there are no remaining free ports.
+func (r *rangeAllocator) fillPorts() {
 	for {
-		port := r.nextFreePort()
-		if port == -1 {
+		if !r.fillPortsOnce() {
 			return
 		}
-		select {
-		case <-stopCh:
-			return
-		case r.ports <- port:
 	}
+}
+
+func (r *rangeAllocator) fillPortsOnce() bool {
+	port := r.nextFreePort()
+	if port == -1 {
+		return false
 	}
+	r.ports <- port
+	return true
 }
 
 // nextFreePort finds a free port, first picking a random port. if that port is already in use

pkg/proxy/winuserspace/port_allocator_test.go

@@ -31,13 +31,26 @@ func TestRangeAllocatorEmpty(t *testing.T) {
 			t.Fatalf("expected panic because of empty port range: %#v", r)
 		}
 	}()
-	_ = newPortRangeAllocator(*r)
+	_ = newPortRangeAllocator(*r, true)
 }
 
 func TestRangeAllocatorFullyAllocated(t *testing.T) {
 	r := &net.PortRange{}
 	r.Set("1-1")
-	a := newPortRangeAllocator(*r)
+	// Don't auto-fill ports, we'll manually turn the crank
+	pra := newPortRangeAllocator(*r, false)
+	a := pra.(*rangeAllocator)
+
+	// Fill in the one available port
+	if !a.fillPortsOnce() {
+		t.Fatalf("Expected to be able to fill ports")
+	}
+
+	// There should be no ports available
+	if a.fillPortsOnce() {
+		t.Fatalf("Expected to be unable to fill ports")
+	}
+
 	p, err := a.AllocateNext()
 	if err != nil {
 		t.Fatalf("unexpected error: %v", err)
@@ -46,12 +59,31 @@ func TestRangeAllocatorFullyAllocated(t *testing.T) {
 		t.Fatalf("unexpected allocated port: %d", p)
 	}
 
+	a.lock.Lock()
+	if bit := a.used.Bit(p - a.Base); bit != 1 {
+		a.lock.Unlock()
+		t.Fatalf("unexpected used bit for allocated port: %d", p)
+	}
+	a.lock.Unlock()
+
 	_, err = a.AllocateNext()
 	if err == nil {
 		t.Fatalf("expected error because of fully-allocated range")
 	}
 
 	a.Release(p)
+	a.lock.Lock()
+	if bit := a.used.Bit(p - a.Base); bit != 0 {
+		a.lock.Unlock()
+		t.Fatalf("unexpected used bit for allocated port: %d", p)
+	}
+	a.lock.Unlock()
+
+	// Fill in the one available port
+	if !a.fillPortsOnce() {
+		t.Fatalf("Expected to be able to fill ports")
+	}
+
 	p, err = a.AllocateNext()
 	if err != nil {
 		t.Fatalf("unexpected error: %v", err)
@@ -59,6 +91,12 @@ func TestRangeAllocatorFullyAllocated(t *testing.T) {
 	if p != 1 {
 		t.Fatalf("unexpected allocated port: %d", p)
 	}
+	a.lock.Lock()
+	if bit := a.used.Bit(p - a.Base); bit != 1 {
+		a.lock.Unlock()
+		t.Fatalf("unexpected used bit for allocated port: %d", p)
+	}
+	a.lock.Unlock()
 
 	_, err = a.AllocateNext()
 	if err == nil {
@@ -69,30 +107,69 @@ func TestRangeAllocatorFullyAllocated(t *testing.T) {
 func TestRangeAllocator_RandomishAllocation(t *testing.T) {
 	r := &net.PortRange{}
 	r.Set("1-100")
-	a := newPortRangeAllocator(*r)
+	pra := newPortRangeAllocator(*r, false)
+	a := pra.(*rangeAllocator)
 
 	// allocate all the ports
 	var err error
 	ports := make([]int, 100, 100)
 	for i := 0; i < 100; i++ {
+		if !a.fillPortsOnce() {
+			t.Fatalf("Expected to be able to fill ports")
+		}
 		ports[i], err = a.AllocateNext()
 		if err != nil {
 			t.Fatalf("unexpected error: %v", err)
 		}
+		if ports[i] < 1 || ports[i] > 100 {
+			t.Fatalf("unexpected allocated port: %d", ports[i])
+		}
+		a.lock.Lock()
+		if bit := a.used.Bit(ports[i] - a.Base); bit != 1 {
+			a.lock.Unlock()
+			t.Fatalf("unexpected used bit for allocated port: %d", ports[i])
+		}
+		a.lock.Unlock()
+	}
+
+	if a.fillPortsOnce() {
+		t.Fatalf("Expected to be unable to fill ports")
 	}
 
 	// release them all
 	for i := 0; i < 100; i++ {
 		a.Release(ports[i])
+		a.lock.Lock()
+		if bit := a.used.Bit(ports[i] - a.Base); bit != 0 {
+			a.lock.Unlock()
+			t.Fatalf("unexpected used bit for allocated port: %d", ports[i])
+		}
+		a.lock.Unlock()
 	}
 
 	// allocate the ports again
 	rports := make([]int, 100, 100)
 	for i := 0; i < 100; i++ {
+		if !a.fillPortsOnce() {
+			t.Fatalf("Expected to be able to fill ports")
+		}
 		rports[i], err = a.AllocateNext()
 		if err != nil {
 			t.Fatalf("unexpected error: %v", err)
 		}
+		if rports[i] < 1 || rports[i] > 100 {
+			t.Fatalf("unexpected allocated port: %d", rports[i])
+		}
+		a.lock.Lock()
+		if bit := a.used.Bit(rports[i] - a.Base); bit != 1 {
+			a.lock.Unlock()
+			t.Fatalf("unexpected used bit for allocated port: %d", rports[i])
+		}
+		a.lock.Unlock()
+	}
+
+	if a.fillPortsOnce() {
+		t.Fatalf("Expected to be unable to fill ports")
 	}
 
 	if reflect.DeepEqual(ports, rports) {