Memory optimization internal blacklist/whitelist cache (#513) (#514)

0xERR0R · May 12, 2022 · 6772438 · 6772438
1 parent 651ab87
commit 6772438
Show file tree

Hide file tree

Showing 5 changed files with 103 additions and 75 deletions.
diff --git a/cache/stringcache/string_caches.go b/cache/stringcache/string_caches.go
@@ -6,8  6,6 @@ import (
 	"strings"
 
 	"github.com/0xERR0R/blocky/log"
-
-	"github.com/0xERR0R/blocky/util"
 )
 
 type StringCache interface {
@@ -22,6  20,10 @@ type CacheFactory interface {
 
 type stringCache map[int]string
 
 func normalizeEntry(entry string) string {
 	return strings.ToLower(entry)
 }
 
 func (cache stringCache) ElementCount() int {
 	count := 0
 
@@ -33,60  35,79 @@ func (cache stringCache) ElementCount() int {
 }
 
 func (cache stringCache) Contains(searchString string) bool {
-	searchLen := len(searchString)
 	normalized := normalizeEntry(searchString)
 	searchLen := len(normalized)
 
 	if searchLen == 0 {
 		return false
 	}
 
 	searchBucketLen := len(cache[searchLen]) / searchLen
 	idx := sort.Search(searchBucketLen, func(i int) bool {
-		return cache[searchLen][i*searchLen:i*searchLen searchLen] >= searchString
 		return cache[searchLen][i*searchLen:i*searchLen searchLen] >= normalized
 	})
 
 	if idx < searchBucketLen {
-		return cache[searchLen][idx*searchLen:idx*searchLen searchLen] == strings.ToLower(searchString)
 		return cache[searchLen][idx*searchLen:idx*searchLen searchLen] == strings.ToLower(normalized)
 	}
 
 	return false
 }
 
 type stringCacheFactory struct {
-	cache stringCache
-	keys  map[string]struct{}
-	tmp   map[int]*strings.Builder
 	// temporary map which holds sorted slice of strings grouped by string length
 	tmp map[int][]string
 }
 
 func newStringCacheFactory() CacheFactory {
 	return &stringCacheFactory{
-		cache: make(stringCache),
-		// temporary map to remove duplicates
-		keys: make(map[string]struct{}),
-		tmp:  make(map[int]*strings.Builder),
 		tmp: make(map[int][]string),
 	}
 }
 
-func (s *stringCacheFactory) AddEntry(entry string) {
-	if _, value := s.keys[entry]; !value {
-		s.keys[entry] = struct{}{}
-		if s.tmp[len(entry)] == nil {
-			s.tmp[len(entry)] = &strings.Builder{}
-		}
 func (s *stringCacheFactory) getBucket(length int) []string {
 	if s.tmp[length] == nil {
 		s.tmp[length] = make([]string, 0)
 	}
 
 	return s.tmp[length]
 }
 
 func (s *stringCacheFactory) insertString(entry string) {
 	normalized := normalizeEntry(entry)
 	entryLen := len(normalized)
 	bucket := s.getBucket(entryLen)
 	ix := sort.SearchStrings(bucket, normalized)
 
-		s.tmp[len(entry)].WriteString(entry)
 	if !(ix < len(bucket) && bucket[ix] == normalized) {
 		// extent internal bucket
 		bucket = append(s.getBucket(entryLen), "")
 
 		// move elements to make place for the insertion
 		copy(bucket[ix 1:], bucket[ix:])
 
 		// insert string at the calculated position
 		bucket[ix] = normalized
 		s.tmp[entryLen] = bucket
 	}
 }
 
 func (s *stringCacheFactory) AddEntry(entry string) {
 	// skip empty strings
 	if len(entry) > 0 {
 		s.insertString(entry)
 	}
 }
 
 func (s *stringCacheFactory) Create() StringCache {
 	cache := make(stringCache, len(s.tmp))
 	for k, v := range s.tmp {
-		chunks := util.Chunks(v.String(), k)
-		sort.Strings(chunks)
-
-		s.cache[k] = strings.Join(chunks, "")
-
-		v.Reset()
 		cache[k] = strings.Join(v, "")
 	}
 
-	return s.cache
 	s.tmp = nil
 
 	return cache
 }
 
 type regexCache []*regexp.Regexp

diff --git a/cache/stringcache/string_caches_benchmark_test.go b/cache/stringcache/string_caches_benchmark_test.go
@@ -0,0  1,44 @@
 package stringcache
 
 import (
 	"math/rand"
 	"testing"
 )
 
 func BenchmarkStringCache(b *testing.B) {
 	testdata := createTestdata(10_000)
 
 	b.ReportAllocs()
 
 	for i := 0; i < b.N; i   {
 		factory := newStringCacheFactory()
 
 		for _, s := range testdata {
 			factory.AddEntry(s)
 		}
 
 		factory.Create()
 	}
 }
 
 func randString(n int) string {
 	const charPool = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789-."
 
 	b := make([]byte, n)
 
 	for i := range b {
 		b[i] = charPool[rand.Intn(len(charPool))] // nolint:gosec
 	}
 
 	return string(b)
 }
 
 func createTestdata(count int) []string {
 	var result []string
 
 	for i := 0; i < count; i   {
 		result = append(result, randString(8 rand.Intn(20))) // nolint:gosec
 	}
 
 	return result
 }
diff --git a/cache/stringcache/string_caches_test.go b/cache/stringcache/string_caches_test.go
@@ -11,13  11,24 @@ var _ = Describe("Caches", func() {
 			factory := newStringCacheFactory()
 			factory.AddEntry("google.com")
 			factory.AddEntry("apple.com")
 			factory.AddEntry("")
 			factory.AddEntry("google.com")
 			factory.AddEntry("APPLe.com")
 
 			cache := factory.Create()
-			It("should match if StringCache Contains string", func() {
 
 			It("should match if StringCache contains exact string", func() {
 				Expect(cache.Contains("apple.com")).Should(BeTrue())
 				Expect(cache.Contains("google.com")).Should(BeTrue())
 				Expect(cache.Contains("www.google.com")).Should(BeFalse())
 				Expect(cache.Contains("")).Should(BeFalse())
 			})
 			It("should match case-insensitive", func() {
 				Expect(cache.Contains("aPPle.com")).Should(BeTrue())
 				Expect(cache.Contains("google.COM")).Should(BeTrue())
 				Expect(cache.Contains("www.google.com")).Should(BeFalse())
 				Expect(cache.Contains("")).Should(BeFalse())
 			})
 			It("should return correct element count", func() {
 				Expect(cache.ElementCount()).Should(Equal(2))
 			})

diff --git a/util/common.go b/util/common.go
@@ -163,34  163,6 @@ func FatalOnError(message string, err error) {
 	}
 }
 
-// Chunks splits the string in multiple chunks
-func Chunks(s string, chunkSize int) []string {
-	if chunkSize >= len(s) {
-		return []string{s}
-	}
-
-	var chunks []string
-
-	chunk := make([]rune, chunkSize)
-	ln := 0
-
-	for _, r := range s {
-		chunk[ln] = r
-		ln  
-
-		if ln == chunkSize {
-			chunks = append(chunks, string(chunk))
-			ln = 0
-		}
-	}
-
-	if ln > 0 {
-		chunks = append(chunks, string(chunk[:ln]))
-	}
-
-	return chunks
-}
-
 // GenerateCacheKey return cacheKey by query type/domain
 func GenerateCacheKey(qType dns.Type, qName string) string {
 	const qTypeLength = 2

diff --git a/util/common_test.go b/util/common_test.go
@@ -16,26  16,6 @@ import (
 )
 
 var _ = Describe("Common function tests", func() {
-	Describe("Split string in chunks", func() {
-		When("String length < chunk size", func() {
-			It("should return one chunk", func() {
-				chunks := Chunks("mystring", 10)
-
-				Expect(chunks).Should(HaveLen(1))
-				Expect(chunks).Should(ContainElement("mystring"))
-			})
-		})
-
-		When("String length > chunk size", func() {
-			It("should return multiple chunks", func() {
-				chunks := Chunks("myveryveryverylongstring", 5)
-
-				Expect(chunks).Should(HaveLen(5))
-				Expect(chunks).Should(ContainElements("myver", "yvery", "veryl", "ongst", "ring"))
-			})
-		})
-	})
-
 	Describe("Print DNS answer", func() {
 		When("different types of DNS answers", func() {
 			rr := make([]dns.RR, 0)