Reduce allocations and time for generator functions

cbrnrd · Dec 10, 2023 · f16752c · f16752c
1 parent fe45dc8
commit f16752c
Show file tree

Hide file tree

Showing 3 changed files with 54 additions and 38 deletions.
diff --git a/README.md b/README.md
@@ -26,39 +26,49 @@ Version: 0.1.0
 
 ```bash
 $ GOMAXPROCS=1 go test -bench=. -benchmem
-goos: darwin
-goarch: arm64
+goos: linux
+goarch: amd64
 pkg: github.com/cbrnrd/ipgen
-BenchmarkGenIPv4               	28962502	        41.33 ns/op	      15 B/op	       1 allocs/op
-BenchmarkGenIPv6               	 6986734	       170.1 ns/op	      63 B/op	       2 allocs/op
-BenchmarkGenIPv4WithExclusions 	 9832515	       120.8 ns/op	      15 B/op	       1 allocs/op
-BenchmarkGenIPv6WithExclusions 	 2491005	       480.8 ns/op	      64 B/op	       2 allocs/op
+cpu: AMD Ryzen 5 2600X Six-Core Processor
+BenchmarkGenIPv4                44527428                26.54 ns/op            4 B/op          1 allocs/op
+BenchmarkGenIPv6                32401519                36.90 ns/op           16 B/op          1 allocs/op
+BenchmarkGenIPv4WithExclusions  21224145                56.56 ns/op            4 B/op          1 allocs/op
+BenchmarkGenIPv6WithExclusions  16818654                71.12 ns/op           16 B/op          1 allocs/op
 ```
 
 
 ### Generating 1,000,000 random IPs
 
+Specs:
+|Component|Spec|
+|---|---|
+|CPU|AMD Ryzen 5 2600X Six-Core Processor|
+|Arch|x86_64|
+|RAM|16GB DDR4|
+|OS|Ubuntu WSL|
+|Go|1.21.4|
+
 ```bash
 $ hyperfine --warmup 3 "./ipgen -n 1000000 -o ips.txt" "./ipgen -n 1000000 -o ips.txt -x 192.168.0.0/16,10.0.0.0/16" "./ipgen -6 -n 1000000 -o ips.txt" "./ipgen -6 -n 1000000 -o ips.txt -x 192.168.0.0/16,10.0.0.0/16"
 Benchmark 1: ./ipgen -n 1000000 -o ips.txt
-  Time (mean ± σ):      2.274 s ±  0.071 s    [User: 1.686 s, System: 2.391 s]
-  Range (min … max):    2.200 s …  2.432 s    10 runs
+  Time (mean ± σ):      3.122 s ±  0.016 s    [User: 1.713 s, System: 3.725 s]
+  Range (min … max):    3.105 s …  3.152 s    10 runs
 
 Benchmark 2: ./ipgen -n 1000000 -o ips.txt -x 192.168.0.0/16,10.0.0.0/16
-  Time (mean ± σ):      2.372 s ±  0.047 s    [User: 1.782 s, System: 2.470 s]
-  Range (min … max):    2.311 s …  2.459 s    10 runs
+  Time (mean ± σ):      3.172 s ±  0.032 s    [User: 1.807 s, System: 3.750 s]
+  Range (min … max):    3.150 s …  3.255 s    10 runs
 
 Benchmark 3: ./ipgen -6 -n 1000000 -o ips.txt
-  Time (mean ± σ):      2.541 s ±  0.061 s    [User: 1.871 s, System: 2.579 s]
-  Range (min … max):    2.466 s …  2.663 s    10 runs
+  Time (mean ± σ):      3.524 s ±  0.103 s    [User: 1.955 s, System: 4.035 s]
+  Range (min … max):    3.374 s …  3.758 s    10 runs
 
 Benchmark 4: ./ipgen -6 -n 1000000 -o ips.txt -x 192.168.0.0/16,10.0.0.0/16
-  Time (mean ± σ):      2.799 s ±  0.053 s    [User: 2.400 s, System: 2.705 s]
-  Range (min … max):    2.724 s …  2.876 s    10 runs
+  Time (mean ± σ):      3.776 s ±  0.137 s    [User: 2.180 s, System: 4.287 s]
+  Range (min … max):    3.519 s …  4.015 s    10 runs
 
 Summary
   ./ipgen -n 1000000 -o ips.txt ran
-    1.04 ± 0.04 times faster than ./ipgen -n 1000000 -o ips.txt -x 192.168.0.0/16,10.0.0.0/16
-    1.12 ± 0.04 times faster than ./ipgen -6 -n 1000000 -o ips.txt
-    1.23 ± 0.04 times faster than ./ipgen -6 -n 1000000 -o ips.txt -x 192.168.0.0/16,10.0.0.0/16
+    1.02 ± 0.01 times faster than ./ipgen -n 1000000 -o ips.txt -x 192.168.0.0/16,10.0.0.0/16
+    1.13 ± 0.03 times faster than ./ipgen -6 -n 1000000 -o ips.txt
+    1.21 ± 0.04 times faster than ./ipgen -6 -n 1000000 -o ips.txt -x 192.168.0.0/16,10.0.0.0/16
 ```
diff --git a/main.go b/main.go
@@ -3,6 +3,7 @@ package main
 import (
 	"flag"
 	"fmt"
+	"io"
 	"net"
 	"os"
 	"strings"
@@ -57,6 +58,7 @@ func main() {
 
 	close(jobs)
 	wg.Wait()
+
 }
 
 func SetupOutput(outpath string) *os.File {
@@ -89,37 +91,38 @@ func ParseExcludedRanges(excludedRangesStr string) []net.IPNet {
 
 // Runs `generatorFunc` and writes the result to `outFile`.
 // Intended to be invoked as a goroutine.
-func Run(wg *sync.WaitGroup, jobs <-chan int, outFile *os.File, generatorFunc func() string) {
+func Run(wg *sync.WaitGroup, jobs <-chan int, out io.Writer, generatorFunc func() net.IP) {
 	defer wg.Done()
 	for range jobs {
-		outFile.WriteString(generatorFunc() + "\n")
+		io.WriteString(out, generatorFunc().String()+"\n")
 	}
 }
 
 // Runs `generatorFunc` and writes the result to `outFile`, excluding any IPs in `excludedRanges`.
 // Intended to be invoked as a goroutine.
-func RunWithExclusions(wg *sync.WaitGroup, jobs <-chan int, outFile *os.File, excludedRanges []net.IPNet, generatorFunc func() string) {
+func RunWithExclusions(wg *sync.WaitGroup, jobs <-chan int, out io.Writer, excludedRanges []net.IPNet, generatorFunc func() net.IP) {
 	defer wg.Done()
 	for range jobs {
 		ip := generatorFunc()
 		if !IsExcluded(ip, excludedRanges) {
-			outFile.WriteString(ip + "\n")
+			io.WriteString(out, ip.String()+"\n")
+
 		}
 	}
 }
 
 // Returns the correct generator function based on the v6 flag
-func GetGenerator(v6 bool, excludedRanges []net.IPNet) func() string {
+func GetGenerator(v6 bool, excludedRanges []net.IPNet) func() net.IP {
 	if v6 {
 		if len(excludedRanges) > 0 {
-			return func() string {
+			return func() net.IP {
 				return GenIPv6WithExclusions(excludedRanges)
 			}
 		}
 		return GenIPv6
 	}
 	if len(excludedRanges) > 0 {
-		return func() string {
+		return func() net.IP {
 			return GenIPv4WithExclusions(excludedRanges)
 		}
 	}

diff --git a/pkg/ip/generators.go b/pkg/ip/generators.go
@@ -9,7 +9,7 @@ import (
 // Generate a random IPv4 address.
 // The IP is guaranteed to be a valid, non loopback address
 // or private address.
-func GenIPv4() string {
+func GenIPv4() net.IP {
 	buf := make([]byte, 4)
 	ip := rand.Uint32()
 	o1, o2 := byte(ip>>24)&0xff, byte(ip>>16)&0xff
@@ -25,38 +25,41 @@ func GenIPv4() string {
 	}
 
 	binary.LittleEndian.PutUint32(buf, ip)
-	return net.IP(buf).String()
+	return net.IP(buf)
 }
 
-// Generate a random IPv6 address.
-func GenIPv6() string {
-	size := 16
-	ip := make([]byte, size)
-	for i := 0; i < size; i++ {
-		ip[i] = byte(rand.Intn(256))
-	}
-	return net.IP(ip).To16().String()
+// Generates a random IPv6 address by creating two random uint64s
+// and shifing them into a byte array
+func GenIPv6() net.IP {
+	buf := make([]byte, 16)
+	ip1 := rand.Uint64()
+	ip2 := rand.Uint64()
+
+	binary.LittleEndian.PutUint64(buf, ip1)
+	binary.LittleEndian.PutUint64(buf[8:], ip2)
+
+	return net.IP(buf)
 }
 
 // Checks if an IP is in any of the excluded ranges
-func IsExcluded(ip string, excludedRanges []net.IPNet) bool {
+func IsExcluded(ip net.IP, excludedRanges []net.IPNet) bool {
 	for _, network := range excludedRanges {
-		if network.Contains(net.ParseIP(ip)) {
+		if network.Contains(ip) {
 			return true
 		}
 	}
 	return false
 }
 
-func GenIPv4WithExclusions(excludedRanges []net.IPNet) string {
+func GenIPv4WithExclusions(excludedRanges []net.IPNet) net.IP {
 	ip := GenIPv4()
 	if IsExcluded(ip, excludedRanges) {
 		return GenIPv4WithExclusions(excludedRanges)
 	}
 	return ip
 }
 
-func GenIPv6WithExclusions(excludedRanges []net.IPNet) string {
+func GenIPv6WithExclusions(excludedRanges []net.IPNet) net.IP {
 	ip := GenIPv6()
 	if IsExcluded(ip, excludedRanges) {
 		return GenIPv6WithExclusions(excludedRanges)