From 9a6788cb137451b8e14d2a24371223f9044373dd Mon Sep 17 00:00:00 2001
From: Kevin Klues <kklues@nvidia.com>
Date: Tue, 13 Aug 2019 17:54:21 +0200
Subject: [PATCH] Add IterateSocketMasks() function to socketmask abstraction

---
 .../topologymanager/socketmask/socketmask.go  | 20 ++++++++
 .../socketmask/socketmask_test.go             | 51 +++++++++++++++++++
 2 files changed, 71 insertions(+)

diff --git a/pkg/kubelet/cm/topologymanager/socketmask/socketmask.go b/pkg/kubelet/cm/topologymanager/socketmask/socketmask.go
index 6a6ce9cc316..becc94eb38e 100644
--- a/pkg/kubelet/cm/topologymanager/socketmask/socketmask.go
+++ b/pkg/kubelet/cm/topologymanager/socketmask/socketmask.go
@@ -185,3 +185,23 @@ func Or(first SocketMask, masks ...SocketMask) SocketMask {
 	s.Or(masks...)
 	return &s
 }
+
+// IterateSocketMasks iterates all possible masks from a list of sockets,
+// issuing a callback on each mask.
+func IterateSocketMasks(sockets []int, callback func(SocketMask)) {
+	var iterate func(sockets, accum []int, size int)
+	iterate = func(sockets, accum []int, size int) {
+		if len(accum) == size {
+			mask, _ := NewSocketMask(accum...)
+			callback(mask)
+			return
+		}
+		for i := range sockets {
+			iterate(sockets[i+1:], append(accum, sockets[i]), size)
+		}
+	}
+
+	for i := 1; i <= len(sockets); i++ {
+		iterate(sockets, []int{}, i)
+	}
+}
diff --git a/pkg/kubelet/cm/topologymanager/socketmask/socketmask_test.go b/pkg/kubelet/cm/topologymanager/socketmask/socketmask_test.go
index eaad0a819e0..76e78bfe16d 100644
--- a/pkg/kubelet/cm/topologymanager/socketmask/socketmask_test.go
+++ b/pkg/kubelet/cm/topologymanager/socketmask/socketmask_test.go
@@ -342,3 +342,54 @@ func TestIsNarrowerThan(t *testing.T) {
 		}
 	}
 }
+
+func TestIterateSocketMasks(t *testing.T) {
+	tcases := []struct {
+		name       string
+		numSockets int
+	}{
+		{
+			name:       "1 Socket",
+			numSockets: 1,
+		},
+		{
+			name:       "2 Sockets",
+			numSockets: 2,
+		},
+		{
+			name:       "4 Sockets",
+			numSockets: 4,
+		},
+		{
+			name:       "8 Sockets",
+			numSockets: 8,
+		},
+		{
+			name:       "16 Sockets",
+			numSockets: 16,
+		},
+	}
+	for _, tc := range tcases {
+		// Generate a list of sockets from tc.numSockets.
+		var sockets []int
+		for i := 0; i < tc.numSockets; i++ {
+			sockets = append(sockets, i)
+		}
+
+		// Calculate the expected number of masks. Since we always have masks
+		// with sockets from 0..n, this is just (2^n - 1) since we want 1 mask
+		// represented by each integer between 1 and 2^n-1.
+		expectedNumMasks := (1 << uint(tc.numSockets)) - 1
+
+		// Iterate all masks and count them.
+		numMasks := 0
+		IterateSocketMasks(sockets, func(SocketMask) {
+			numMasks++
+		})
+
+		// Compare the number of masks generated to the expected amount.
+		if expectedNumMasks != numMasks {
+			t.Errorf("Expected to iterate %v masks, got %v", expectedNumMasks, numMasks)
+		}
+	}
+}