add benchmarks

kubernetes · Apr 14, 2022 · 3ccc71f · 3ccc71f
1 parent d3259b7
commit 3ccc71f
Show file tree

Hide file tree

Showing 5 changed files with 108 additions and 40 deletions.
diff --git a/pkg/net/cidrs/aws/mapper_test.go b/pkg/net/cidrs/aws/mapper_test.go
@@ -19,27 +19,30 @@ package aws
 import (
 	"net/netip"
 	"testing"
+
+	"sigs.k8s.io/oci-proxy/pkg/net/cidrs"
 )
 
+var testCases = []struct {
+	Addr           netip.Addr
+	ExpectedRegion string
+}{
+	// some known IPs and their regions
+	{Addr: netip.MustParseAddr("35.180.1.1"), ExpectedRegion: "eu-west-3"},
+	{Addr: netip.MustParseAddr("35.250.1.1"), ExpectedRegion: ""},
+	{Addr: netip.MustParseAddr("35.0.1.1"), ExpectedRegion: ""},
+	{Addr: netip.MustParseAddr("52.94.76.1"), ExpectedRegion: "us-west-2"},
+	{Addr: netip.MustParseAddr("52.94.77.1"), ExpectedRegion: "us-west-2"},
+	{Addr: netip.MustParseAddr("52.93.127.172"), ExpectedRegion: "us-east-1"},
+	// ipv6
+	{Addr: netip.MustParseAddr("2400:6500:0:9::2"), ExpectedRegion: "ap-southeast-3"},
+	{Addr: netip.MustParseAddr("2400:6500:0:9::1"), ExpectedRegion: "ap-southeast-3"},
+	{Addr: netip.MustParseAddr("2400:6500:0:9::3"), ExpectedRegion: "ap-southeast-3"},
+	{Addr: netip.MustParseAddr("2600:1f01:4874::47"), ExpectedRegion: "us-west-2"},
+	{Addr: netip.MustParseAddr("2400:6500:0:9::100"), ExpectedRegion: ""},
+}
+
 func TestNewAWSRegionMapper(t *testing.T) {
-	testCases := []struct {
-		Addr           netip.Addr
-		ExpectedRegion string
-	}{
-		// some known IPs and their regions
-		{Addr: netip.MustParseAddr("35.180.1.1"), ExpectedRegion: "eu-west-3"},
-		{Addr: netip.MustParseAddr("35.250.1.1"), ExpectedRegion: ""},
-		{Addr: netip.MustParseAddr("35.0.1.1"), ExpectedRegion: ""},
-		{Addr: netip.MustParseAddr("52.94.76.1"), ExpectedRegion: "us-west-2"},
-		{Addr: netip.MustParseAddr("52.94.77.1"), ExpectedRegion: "us-west-2"},
-		{Addr: netip.MustParseAddr("52.93.127.172"), ExpectedRegion: "us-east-1"},
-		// ipv6
-		{Addr: netip.MustParseAddr("2400:6500:0:9::2"), ExpectedRegion: "ap-southeast-3"},
-		{Addr: netip.MustParseAddr("2400:6500:0:9::1"), ExpectedRegion: "ap-southeast-3"},
-		{Addr: netip.MustParseAddr("2400:6500:0:9::3"), ExpectedRegion: "ap-southeast-3"},
-		{Addr: netip.MustParseAddr("2600:1f01:4874::47"), ExpectedRegion: "us-west-2"},
-		{Addr: netip.MustParseAddr("2400:6500:0:9::100"), ExpectedRegion: ""},
-	}
 	mapper := NewAWSRegionMapper()
 	for i := range testCases {
 		tc := testCases[i]
@@ -55,3 +58,53 @@ func TestNewAWSRegionMapper(t *testing.T) {
 		})
 	}
 }
+
+/*  for benchmarking memory / init time */
+
+func BenchmarkNewAWSRegionMapper(b *testing.B) {
+	for n := 0; n < b.N; n++ {
+		mapper := NewAWSRegionMapper()
+		// get any address just to prevent mapper being optimized out
+		mapper.GetIP(testCases[0].Addr)
+	}
+}
+
+func BenchmarkNewAWSegionBruteForce(b *testing.B) {
+	for n := 0; n < b.N; n++ {
+		mapper := cidrs.NewBruteForceMapper(regionToRanges)
+		// get any address just to prevent mapper being optimized out
+		mapper.GetIP(testCases[0].Addr)
+	}
+}
+
+/* for benchmarking matching time */
+
+func BenchmarkAWSRegionTrieMap(b *testing.B) {
+	mapper := NewAWSRegionMapper()
+	for n := 0; n < b.N; n++ {
+		tc := testCases[n%len(testCases)]
+		region, matched := mapper.GetIP(tc.Addr)
+		expectMatched := tc.ExpectedRegion != ""
+		if matched != expectMatched || region != tc.ExpectedRegion {
+			b.Fatalf(
+				"result does not match for %v, got: (%q, %t) expected: (%q, %t)",
+				tc.Addr, region, matched, tc.ExpectedRegion, expectMatched,
+			)
+		}
+	}
+}
+
+func BenchmarkAWSRegionBruteForce(b *testing.B) {
+	mapper := cidrs.NewBruteForceMapper(regionToRanges)
+	for n := 0; n < b.N; n++ {
+		tc := testCases[n%len(testCases)]
+		region, matched := mapper.GetIP(tc.Addr)
+		expectMatched := tc.ExpectedRegion != ""
+		if matched != expectMatched || region != tc.ExpectedRegion {
+			b.Fatalf(
+				"result does not match for %v, got: (%q, %t) expected: (%q, %t)",
+				tc.Addr, region, matched, tc.ExpectedRegion, expectMatched,
+			)
+		}
+	}
+}
diff --git a/pkg/net/cidrs/bruteforce.go b/pkg/net/cidrs/bruteforce.go
@@ -24,20 +24,22 @@ import (
 // compares all netip.Prefix
 //
 // This type exists purely for testing and benchmarking
-func NewBruteForceMapper[V comparable](mapping map[netip.Prefix]V) IPMapper[V] {
+func NewBruteForceMapper[V comparable](mapping map[V][]netip.Prefix) IPMapper[V] {
 	return &bruteForceMapper[V]{
 		mapping: mapping,
 	}
 }
 
 type bruteForceMapper[V comparable] struct {
-	mapping map[netip.Prefix]V
+	mapping map[V][]netip.Prefix
 }
 
 func (b *bruteForceMapper[V]) GetIP(addr netip.Addr) (value V, matched bool) {
-	for cidr, v := range b.mapping {
-		if cidr.Contains(addr) {
-			return v, true
+	for v, cidrs := range b.mapping {
+		for _, cidr := range cidrs {
+			if cidr.Contains(addr) {
+				return v, true
+			}
 		}
 	}
 	return

diff --git a/pkg/net/cidrs/testdatacommon_test.go b/pkg/net/cidrs/testdatacommon_test.go
@@ -21,21 +21,30 @@ import (
 )
 
 // common test data
-var testCIDRS = map[netip.Prefix]string{
-	netip.MustParsePrefix("35.180.0.0/16"):    "eu-west-3",
-	netip.MustParsePrefix("52.94.76.0/22"):    "us-west-2",
-	netip.MustParsePrefix("52.93.127.17/32"):  "eu-west-3",
-	netip.MustParsePrefix("52.93.127.172/31"): "eu-west-3",
-	netip.MustParsePrefix("52.93.127.173/32"): "us-east-1",
-	netip.MustParsePrefix("52.93.127.174/32"): "ap-northeast-1",
-	netip.MustParsePrefix("52.93.127.175/32"): "ap-northeast-1",
-	netip.MustParsePrefix("52.93.127.176/32"): "ap-northeast-1",
-	netip.MustParsePrefix("52.93.127.177/32"): "ap-northeast-1",
-	netip.MustParsePrefix("52.93.127.178/32"): "ap-northeast-1",
-	netip.MustParsePrefix("52.93.127.179/32"): "ap-northeast-1",
-	// ipv6
-	netip.MustParsePrefix("2400:6500:0:9::2/128"): "ap-southeast-3",
-	netip.MustParsePrefix("2600:1f01:4874::/47"):  "us-west-2",
+var testCIDRS = map[string][]netip.Prefix{
+	"eu-west-3": {
+		netip.MustParsePrefix("35.180.0.0/16"),
+		netip.MustParsePrefix("52.93.127.17/32"),
+		netip.MustParsePrefix("52.93.127.172/31"),
+	},
+	"us-east-1": {
+		netip.MustParsePrefix("52.93.127.173/32"),
+	},
+	"us-west-2": {
+		netip.MustParsePrefix("2600:1f01:4874::/47"),
+		netip.MustParsePrefix("52.94.76.0/22"),
+	},
+	"ap-northeast-1": {
+		netip.MustParsePrefix("52.93.127.174/32"),
+		netip.MustParsePrefix("52.93.127.175/32"),
+		netip.MustParsePrefix("52.93.127.176/32"),
+		netip.MustParsePrefix("52.93.127.177/32"),
+		netip.MustParsePrefix("52.93.127.178/32"),
+		netip.MustParsePrefix("52.93.127.179/32"),
+	},
+	"ap-southeast-3": {
+		netip.MustParsePrefix("2400:6500:0:9::2/128"),
+	},
 }
 
 // common test cases

diff --git a/pkg/net/cidrs/triemap.go b/pkg/net/cidrs/triemap.go
@@ -29,6 +29,8 @@ import "net/netip"
 // Currently this is a simple TrieMap, in the future it may have compression.
 //
 // See: https://vincent.bernat.ch/en/blog/2017-ipv4-route-lookup-linux
+//
+// For benchmarks with real data see ./aws/mapper_test.go
 type TrieMap[V comparable] struct {
 	// This is the real triemap, but it only maps netip.Prefix / netip.Addr : int
 	// see: https://planetscale.com/blog/generics-can-make-your-go-code-slower

diff --git a/pkg/net/cidrs/triemap_test.go b/pkg/net/cidrs/triemap_test.go
@@ -22,8 +22,10 @@ import (
 
 func TestTrieMap(t *testing.T) {
 	trieMap := NewTrieMap[string]()
-	for cidr, value := range testCIDRS {
-		trieMap.Insert(cidr, value)
+	for value, cidrs := range testCIDRS {
+		for _, cidr := range cidrs {
+			trieMap.Insert(cidr, value)
+		}
 	}
 	for i := range testCases {
 		tc := testCases[i]