Repository: DiceDB/dice
File: core/queueref_test.go (Go)

Code before:

// Test for removing from queue with all expired keys
func TestRemoveAllExpired(t *testing.T) {
	qr := core.NewQueueRef()
	val := [3]int{10, 20, 30}
	core.Put("key1", core.NewObj(val[0], 1, core.OBJ_TYPE_STRING, core.OBJ_ENCODING_INT))
	qr.Insert("key1")
	core.Put("key2", core.NewObj(val[1], 1, core.OBJ_TYPE_STRING, core.OBJ_ENCODING_INT))
	qr.Insert("key2")
	core.Put("key3", core.NewObj(val[2], 1, core.OBJ_TYPE_STRING, core.OBJ_ENCODING_INT))
	qr.Insert("key3")
	time.Sleep(2 * time.Millisecond)
	// remove from empty queue
	_, err := qr.Remove()
	assert.Equal(t, err, core.ErrQueueEmpty, fmt.Sprintf("test for removing from empty queue failed Expected : %s, Got : %s\n", core.ErrQueueEmpty, err))
}

// Benchmark queueref by inserting expired, non-expired and expired keys in order and removing them
func BenchmarkQueueRef(b *testing.B) {
	benchmarkCases := []struct {
		name         string
		expiredFirst int
		nonExpired   int
		expiredLast  int
	}{
		{"Small", 10000, 20000, 30000},
		{"Medium", 100000, 200000, 300000},
		{"Large", 1000000, 2000000, 3000000},
		{"VeryLarge", 10000000, 20000000, 30000000},
	}

	// Setup code...
	qr := core.NewQueueRef()
	config.KeysLimit = 100000000 // override keys limit high for benchmarking

	// need to init again for each round with the overriden buffer size
	// otherwise the watchchannel buffer size will stay as it is with the global keylimits size
	core.WatchChannel = make(chan core.WatchEvent, config.KeysLimit)
	b.ResetTimer()

	for _, bc := range benchmarkCases {
		b.Run(fmt.Sprintf("%s Insert expired", bc.name), func(b *testing.B) {
			for i := 0; i < bc.expiredFirst; i++ {
				// Insertion benchmark...
				key := fmt.Sprintf("k%d", i)
				value := fmt.Sprintf("v%d", i)
				core.Put(key, core.NewObj(value, 1, core.OBJ_TYPE_STRING, core.OBJ_ENCODING_INT))
				qr.Insert(key)
			}
		})


Review comment (suggestion):
Trying to understand something. Should the `expiredLast` not be equal to `nonExpired`?
You are waiting for 2ms before you run the benchmark for qr.Remove(), wouldn't the `expiredFirst` keys have expired by then?

If my understanding is correct then it may be better to call the Remove() operation only `nonExpired` number of times, and finally verify that the queue is empty at the end of it (you can call b.StopTimer() before the assertion).

Almost maybe it would be better to rename the fields to expiredKeys, nonExpiredKeys.

Code after:
}

// Test for removing from queue with all expired keys
func TestRemoveAllExpired(t *testing.T) {
	qr := core.NewQueueRef()
	val := [3]int{10, 20, 30}
	core.Put("key1", core.NewObj(val[0], 1, core.OBJ_TYPE_STRING, core.OBJ_ENCODING_INT))
	qr.Insert("key1")
	core.Put("key2", core.NewObj(val[1], 1, core.OBJ_TYPE_STRING, core.OBJ_ENCODING_INT))
	qr.Insert("key2")
	core.Put("key3", core.NewObj(val[2], 1, core.OBJ_TYPE_STRING, core.OBJ_ENCODING_INT))
	qr.Insert("key3")
	time.Sleep(2 * time.Millisecond)
	// remove from empty queue
	_, err := qr.Remove()
	assert.Equal(t, err, core.ErrQueueEmpty, fmt.Sprintf("test for removing from empty queue failed Expected : %s, Got : %s\n", core.ErrQueueEmpty, err))
}

// Benchmark queueref by inserting expired, non-expired and expired keys in order and removing them
func BenchmarkQueueRef(b *testing.B) {
	// Setup code...
	qr := core.NewQueueRef()
	config.KeysLimit = 100000000 // override keys limit high for benchmarking

	// need to init again for each round with the overriden buffer size
	// otherwise the watchchannel buffer size will stay as it is with the global keylimits size
	core.WatchChannel = make(chan core.WatchEvent, config.KeysLimit)

	expiredCount := 0
	nonExpiredCount := 0
	b.ResetTimer()
	b.Run("Insertion of expired and non-expired keys", func(b *testing.B) {
		for i := 0; i < b.N; i++ {
			key := fmt.Sprintf("k%d", i)
			value := fmt.Sprintf("v%d", i)
			random := rand.Intn(2)
			// Insert keys without expiration
			expiration := -1
			if random == 1 {
				// Insert keys with expiration of 1ms
				expiration = 1
				expiredCount++
			} else {
				nonExpiredCount++
			}
			core.Put(key, core.NewObj(value, int64(expiration), core.OBJ_TYPE_STRING, core.OBJ_ENCODING_INT))
			qr.Insert(key)
		}
	})
	b.StopTimer()
	// Allow keys to expire