Merge pull request #3141 from corhere/local-idm

Migrate to local version of package idm
moby · Jul 6, 2023 · 957f1e9 · 957f1e9
2 parents 93fe90a + 64a3444
commit 957f1e9
Show file tree

Hide file tree

Showing 3 changed files with 358 additions and 10 deletions.
diff --git a/internal/idm/idm.go b/internal/idm/idm.go
@@ -0,0 +1,69 @@
+// Package idm manages reservation/release of numerical ids from a configured set of contiguous ids.
+package idm
+
+import (
+	"errors"
+	"fmt"
+
+	"github.com/docker/docker/libnetwork/bitmap"
+)
+
+// IDM manages the reservation/release of numerical ids from a contiguous set.
+//
+// An IDM instance is not safe for concurrent use.
+type IDM struct {
+	start  uint64
+	end    uint64
+	handle *bitmap.Bitmap
+}
+
+// New returns an instance of id manager for a [start,end] set of numerical ids.
+func New(start, end uint64) (*IDM, error) {
+	if end <= start {
+		return nil, fmt.Errorf("invalid set range: [%d, %d]", start, end)
+	}
+
+	return &IDM{start: start, end: end, handle: bitmap.New(1 + end - start)}, nil
+}
+
+// GetID returns the first available id in the set.
+func (i *IDM) GetID(serial bool) (uint64, error) {
+	if i.handle == nil {
+		return 0, errors.New("ID set is not initialized")
+	}
+	ordinal, err := i.handle.SetAny(serial)
+	return i.start + ordinal, err
+}
+
+// GetSpecificID tries to reserve the specified id.
+func (i *IDM) GetSpecificID(id uint64) error {
+	if i.handle == nil {
+		return errors.New("ID set is not initialized")
+	}
+
+	if id < i.start || id > i.end {
+		return errors.New("requested id does not belong to the set")
+	}
+
+	return i.handle.Set(id - i.start)
+}
+
+// GetIDInRange returns the first available id in the set within a [start,end] range.
+func (i *IDM) GetIDInRange(start, end uint64, serial bool) (uint64, error) {
+	if i.handle == nil {
+		return 0, errors.New("ID set is not initialized")
+	}
+
+	if start < i.start || end > i.end {
+		return 0, errors.New("requested range does not belong to the set")
+	}
+
+	ordinal, err := i.handle.SetAnyInRange(start-i.start, end-i.start, serial)
+
+	return i.start + ordinal, err
+}
+
+// Release releases the specified id.
+func (i *IDM) Release(id uint64) {
+	i.handle.Unset(id - i.start)
+}
diff --git a/internal/idm/idm_test.go b/internal/idm/idm_test.go
@@ -0,0 +1,284 @@
+package idm
+
+import (
+	"testing"
+)
+
+func TestNew(t *testing.T) {
+	i, err := New(0, 10)
+	if err != nil {
+		t.Errorf("idm.New(0, 10) error = %v", err)
+	}
+	if i.handle == nil {
+		t.Error("set is not initialized")
+	}
+	if i.start != 0 {
+		t.Errorf("unexpected start: got %d, want 0", i.start)
+	}
+	if i.end != 10 {
+		t.Errorf("unexpected end: got %d, want 10", i.end)
+	}
+}
+
+func TestAllocate(t *testing.T) {
+	i, err := New(50, 52)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if err = i.GetSpecificID(49); err == nil {
+		t.Error("i.GetSpecificID(49): expected failure but succeeded")
+	}
+
+	if err = i.GetSpecificID(53); err == nil {
+		t.Fatal("i.GetSpecificID(53): expected failure but succeeded")
+	}
+
+	o, err := i.GetID(false)
+	if err != nil {
+		t.Errorf("i.GetID(false) error = %v", err)
+	}
+	if o != 50 {
+		t.Errorf("i.GetID(false) = %v, want 50", o)
+	}
+
+	err = i.GetSpecificID(50)
+	if err == nil {
+		t.Error("i.GetSpecificID(50): allocating already-allocated id should fail")
+	}
+
+	o, err = i.GetID(false)
+	if err != nil {
+		t.Errorf("i.GetID(false) error = %v", err)
+	}
+	if o != 51 {
+		t.Errorf("i.GetID(false) = %v, want 51", o)
+	}
+
+	o, err = i.GetID(false)
+	if err != nil {
+		t.Errorf("i.GetID(false) error = %v", err)
+	}
+	if o != 52 {
+		t.Errorf("i.GetID(false) = %v, want 51", o)
+	}
+
+	o, err = i.GetID(false)
+	if err == nil {
+		t.Errorf("i.GetID(false) = %v, allocating ID from full set should fail", o)
+	}
+
+	i.Release(50)
+
+	o, err = i.GetID(false)
+	if err != nil {
+		t.Errorf("i.GetID(false) error = %v", err)
+	}
+	if o != 50 {
+		t.Errorf("i.GetID(false) = %v, want 50", o)
+	}
+
+	i.Release(52)
+	err = i.GetSpecificID(52)
+	if err != nil {
+		t.Errorf("i.GetSpecificID(52) error = %v, expected success allocating a released ID", err)
+	}
+}
+
+func TestUninitialized(t *testing.T) {
+	i := &IDM{}
+
+	if _, err := i.GetID(false); err == nil {
+		t.Error("i.GetID(...) on uninitialized set should fail")
+	}
+
+	if err := i.GetSpecificID(44); err == nil {
+		t.Error("i.GetSpecificID(...) on uninitialized set should fail")
+	}
+}
+
+func TestAllocateInRange(t *testing.T) {
+	i, err := New(5, 10)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	o, err := i.GetIDInRange(6, 6, false)
+	if err != nil {
+		t.Errorf("i.GetIDInRange(6, 6, false) error = %v", err)
+	}
+	if o != 6 {
+		t.Errorf("i.GetIDInRange(6, 6, false) = %d, want 6", o)
+	}
+
+	if err = i.GetSpecificID(6); err == nil {
+		t.Errorf("i.GetSpecificID(6): allocating already-allocated id should fail")
+	}
+
+	o, err = i.GetID(false)
+	if err != nil {
+		t.Errorf("i.GetID(false) error = %v", err)
+	}
+	if o != 5 {
+		t.Errorf("i.GetID(false) = %v, want 5", o)
+	}
+
+	i.Release(6)
+
+	o, err = i.GetID(false)
+	if err != nil {
+		t.Errorf("i.GetID(false) error = %v", err)
+	}
+	if o != 6 {
+		t.Errorf("i.GetID(false) = %v, want 6", o)
+	}
+
+	for n := uint64(7); n <= 10; n++ {
+		o, err := i.GetIDInRange(7, 10, false)
+		if err != nil {
+			t.Errorf("i.GetIDInRange(7, 10, false) error = %v", err)
+		}
+		if o != n {
+			t.Errorf("i.GetIDInRange(7, 10, false) = %d, want %d", o, n)
+		}
+	}
+
+	if err = i.GetSpecificID(7); err == nil {
+		t.Errorf("i.GetSpecificID(7): allocating already-allocated id should fail")
+	}
+
+	if err = i.GetSpecificID(10); err == nil {
+		t.Errorf("i.GetSpecificID(10): allocating already-allocated id should fail")
+	}
+
+	i.Release(10)
+
+	o, err = i.GetIDInRange(5, 10, false)
+	if err != nil {
+		t.Errorf("i.GetIDInRange(5, 10, false) error = %v", err)
+	}
+	if o != 10 {
+		t.Errorf("i.GetIDInRange(5, 10, false) = %d, want 10", o)
+	}
+
+	i.Release(5)
+
+	o, err = i.GetIDInRange(5, 10, false)
+	if err != nil {
+		t.Errorf("i.GetIDInRange(5, 10, false) error = %v", err)
+	}
+	if o != 5 {
+		t.Errorf("i.GetIDInRange(5, 10, false) = %d, want 5", o)
+	}
+
+	for n := uint64(5); n <= 10; n++ {
+		i.Release(n)
+	}
+
+	for n := uint64(5); n <= 10; n++ {
+		o, err := i.GetIDInRange(5, 10, false)
+		if err != nil {
+			t.Errorf("i.GetIDInRange(5, 10, false) error = %v", err)
+		}
+		if o != n {
+			t.Errorf("i.GetIDInRange(5, 10, false) = %d, want %d", o, n)
+		}
+	}
+
+	for n := uint64(5); n <= 10; n++ {
+		if err = i.GetSpecificID(n); err == nil {
+			t.Errorf("i.GetSpecificID(%d): allocating already-allocated id should fail", n)
+		}
+	}
+
+	// New larger set
+	const ul = (1 << 24) - 1
+	i, err = New(0, ul)
+	if err != nil {
+		t.Fatalf("New(0, %d) error = %v", ul, err)
+	}
+
+	o, err = i.GetIDInRange(4096, ul, false)
+	if err != nil {
+		t.Errorf("i.GetIDInRange(4096, %d, false) error = %v", ul, err)
+	}
+	if o != 4096 {
+		t.Errorf("i.GetIDInRange(4096, %d, false) = %d, want 4096", ul, o)
+	}
+
+	o, err = i.GetIDInRange(4096, ul, false)
+	if err != nil {
+		t.Errorf("i.GetIDInRange(4096, %d, false) error = %v", ul, err)
+	}
+	if o != 4097 {
+		t.Errorf("i.GetIDInRange(4096, %d, false) = %d, want 4097", ul, o)
+	}
+
+	o, err = i.GetIDInRange(4096, ul, false)
+	if err != nil {
+		t.Errorf("i.GetIDInRange(4096, %d, false) error = %v", ul, err)
+	}
+	if o != 4098 {
+		t.Errorf("i.GetIDInRange(4096, %d, false) = %d, want 4098", ul, o)
+	}
+}
+
+func TestAllocateSerial(t *testing.T) {
+	i, err := New(50, 55)
+	if err != nil {
+		t.Fatalf("New(50, 55) error = %v", err)
+	}
+
+	if err = i.GetSpecificID(49); err == nil {
+		t.Errorf("i.GetSpecificID(49): allocating out-of-range id should fail")
+	}
+
+	if err = i.GetSpecificID(56); err == nil {
+		t.Errorf("i.GetSpecificID(56): allocating out-of-range id should fail")
+	}
+
+	o, err := i.GetID(true)
+	if err != nil {
+		t.Errorf("i.GetID(true) error = %v", err)
+	}
+	if o != 50 {
+		t.Errorf("i.GetID(true) = %v, want 50", o)
+	}
+
+	err = i.GetSpecificID(50)
+	if err == nil {
+		t.Errorf("i.GetSpecificID(50): allocating already-allocated id should fail")
+	}
+
+	o, err = i.GetID(true)
+	if err != nil {
+		t.Errorf("i.GetID(true) error = %v", err)
+	}
+	if o != 51 {
+		t.Errorf("i.GetID(true) = %v, want 51", o)
+	}
+
+	o, err = i.GetID(true)
+	if err != nil {
+		t.Errorf("i.GetID(true) error = %v", err)
+	}
+	if o != 52 {
+		t.Errorf("i.GetID(true) = %v, want 52", o)
+	}
+
+	i.Release(50)
+
+	o, err = i.GetID(true)
+	if err != nil {
+		t.Errorf("i.GetID(true) error = %v", err)
+	}
+	if o != 53 {
+		t.Errorf("i.GetID(true) = %v, want 53", o)
+	}
+
+	i.Release(52)
+	err = i.GetSpecificID(52)
+	if err != nil {
+		t.Errorf("i.GetSpecificID(52) error = %v, expected success allocating a released ID", err)
+	}
+}
diff --git a/manager/allocator/cnmallocator/portallocator.go b/manager/allocator/cnmallocator/portallocator.go
@@ -1,10 +1,8 @@
 package cnmallocator
 
 import (
-	"fmt"
-
-	"github.com/docker/docker/libnetwork/idm"
 	"github.com/moby/swarmkit/v2/api"
+	"github.com/moby/swarmkit/v2/internal/idm"
 )
 
 const (
@@ -34,8 +32,8 @@ type portAllocator struct {
 
 type portSpace struct {
 	protocol         api.PortConfig_Protocol
-	masterPortSpace  *idm.Idm
-	dynamicPortSpace *idm.Idm
+	masterPortSpace  *idm.IDM
+	dynamicPortSpace *idm.IDM
 }
 
 type allocatedPorts map[api.PortConfig]map[uint32]*api.PortConfig
@@ -118,15 +116,12 @@ func newPortAllocator() (*portAllocator, error) {
 }
 
 func newPortSpace(protocol api.PortConfig_Protocol) (*portSpace, error) {
-	masterName := fmt.Sprintf("%s-master-ports", protocol)
-	dynamicName := fmt.Sprintf("%s-dynamic-ports", protocol)
-
-	master, err := idm.New(nil, masterName, masterPortStart, masterPortEnd)
+	master, err := idm.New(masterPortStart, masterPortEnd)
 	if err != nil {
 		return nil, err
 	}
 
-	dynamic, err := idm.New(nil, dynamicName, dynamicPortStart, dynamicPortEnd)
+	dynamic, err := idm.New(dynamicPortStart, dynamicPortEnd)
 	if err != nil {
 		return nil, err
 	}