Add a optimized database

pkg/compliance/db.go

@@ -2,74 +2,85 @@ package compliance
 
 import (
 	"fmt"
-
-	"github.com/samber/lo"
 )
 
 type db struct {
-	records []*record
+	keyRecords   map[int][]*record            // store record as a value of a Map with a key as a "check_key"
+	idRecords    map[string][]*record         // store record as a value of a Map with a key as a "id"
+	idKeyRecords map[string]map[int][]*record // store record as a value of a Map with a key as a "check_key an id"
+	allIds       map[string]struct{}          // Set of all unique ids
 }
 
+// newDB initializes and returns a new database instance.
 func newDB() *db {
-	return &db{}
+	return &db{
+		keyRecords:   make(map[int][]*record),
+		idRecords:    make(map[string][]*record),
+		idKeyRecords: make(map[string]map[int][]*record),
+		allIds:       make(map[string]struct{}),
+	}
 }
 
+// addRecord adds a single record to the database
 func (d *db) addRecord(r *record) {
-	d.records = append(d.records, r)
+	// store record using a key
+	d.keyRecords[r.check_key] = append(d.keyRecords[r.check_key], r)
+
+	// store record using a id
+	d.idRecords[r.id] = append(d.idRecords[r.id], r)
+	if d.idKeyRecords[r.id] == nil {
+		d.idKeyRecords[r.id] = make(map[int][]*record)
+	}
+
+	// store record using a key and id
+	d.idKeyRecords[r.id][r.check_key] = append(d.idKeyRecords[r.id][r.check_key], r)
+
+	d.allIds[r.id] = struct{}{}
 }
 
+// addRecords adds multiple records to the database
 func (d *db) addRecords(rs []*record) {
-	d.records = append(d.records, rs...)
+	for _, r := range rs {
+		d.addRecord(r)
+	}
 }
 
-func (d *db) getRecords(key int) []record {
-	var rs []record
-	for _, r := range d.records {
-		if r.check_key == key {
-			rs = append(rs, *r)
-		}
-	}
-	return rs
+// getRecords retrieves records by the given "check_key"
+func (d *db) getRecords(key int) []*record {
+	return d.keyRecords[key]
 }
 
+// getAllIds retrieves all unique ids in the database
 func (d *db) getAllIds() []string {
-	var ids []string
-	for _, r := range d.records {
-		ids = append(ids, r.id)
+	ids := make([]string, 0, len(d.allIds))
+	for id := range d.allIds {
+		ids = append(ids, id)
 	}
-
-	return lo.Uniq(ids)
+	return ids
 }
 
-func (d *db) getRecordsById(id string) []record {
-	var rs []record
-	for _, r := range d.records {
-		if r.id == id {
-			rs = append(rs, *r)
-		}
-	}
-	return rs
+// getRecordsById retrieves records by the given "id"
+func (d *db) getRecordsById(id string) []*record {
+	return d.idRecords[id]
 }
 
-func (d *db) getRecordsByKeyId(key int, id string) []record {
-	var rs []record
-	for _, r := range d.records {
-		if r.check_key == key && r.id == id {
-			rs = append(rs, *r)
-		}
-	}
-	return rs
+// getRecordsByKeyId retrieves records by the given "check_key" and "id"
+func (d *db) getRecordsByKeyId(key int, id string) []*record {
+	return d.idKeyRecords[id][key]
 }
 
-func (d *db) dumpAll(key []int) {
-	for _, r := range d.records {
-		if len(key) == 0 {
-			fmt.Printf("id: %s, key: %d, value: %s\n", r.id, r.check_key, r.check_value)
-			continue
-		}
-		for _, k := range key {
-			if r.check_key == k {
+// dumpAll prints all records, optionally filtered by the given keys
+func (d *db) dumpAll(keys []int) {
+	for _, records := range d.keyRecords {
+		for _, r := range records {
+			if len(keys) == 0 {
 				fmt.Printf("id: %s, key: %d, value: %s\n", r.id, r.check_key, r.check_value)
+				continue
+			}
+			for _, k := range keys {
+				if r.check_key == k {
+					fmt.Printf("id: %s, key: %d, value: %s\n", r.id, r.check_key, r.check_value)
+				}
 			}
 		}
 	}

pkg/compliance/db_test.go

@@ -0,0 +1,69 @@
+package compliance
+
+import (
+	"fmt"
+	"math/rand"
+	"testing"
+)
+
+const (
+	numRecords = 1000000 // Number of records to test with
+)
+
+// Generate large data set
+func generateRecords(n int) []*record {
+	var records []*record
+	for i := 0; i < n; i++ {
+		records = append(records, &record{
+			check_key:   rand.Intn(1000),
+			check_value: fmt.Sprintf("value_%d", i),
+			id:          fmt.Sprintf("id_%d", rand.Intn(1000)),
+			score:       rand.Float64() * 100,
+			required:    rand.Intn(2) == 0,
+		})
+	}
+	return records
+}
+
+// Benchmark original db implementation
+func BenchmarkOriginalDB(b *testing.B) {
+	records := generateRecords(numRecords)
+	db := newDB()
+
+	// Benchmark insertion
+	b.Run("Insert", func(b *testing.B) {
+		for i := 0; i < b.N; i++ {
+			db.addRecords(records)
+		}
+	})
+
+	// Benchmark retrieval by key
+	b.Run("GetByKey", func(b *testing.B) {
+		for i := 0; i < b.N; i++ {
+			db.getRecords(rand.Intn(1000))
+		}
+	})
+
+	// Benchmark retrieval by ID
+	b.Run("GetByID", func(b *testing.B) {
+		for i := 0; i < b.N; i++ {
+			db.getRecordsById(fmt.Sprintf("id_%d", rand.Intn(1000)))
+		}
+	})
+
+	// Benchmark for combined retrieval by key and ID case
+	b.Run("GetByKeyAndIDTogether", func(b *testing.B) {
+		for i := 0; i < b.N; i++ {
+			key := rand.Intn(1000)
+			id := fmt.Sprintf("id_%d", rand.Intn(1000))
+			db.getRecordsByKeyId(key, id)
+		}
+	})
+
+	// Benchmark for retrieval of all IDs case
+	b.Run("GetAllIDs", func(b *testing.B) {
+		for i := 0; i < b.N; i++ {
+			db.getAllIds()
+		}
+	})
+}