Files
cocos/manager/ttl_test.go
T
Sammy Kerata Oina 45187d7f41 COCOS-454 - Implement graceful shutdown for services and add TTL management for VMs (#473)
* Implement graceful shutdown for services and add TTL management for VMs

Signed-off-by: Sammy Oina <sammyoina@gmail.com>

* Remove unnecessary comment from go-tdx-guest dependency in go.mod

Signed-off-by: Sammy Oina <sammyoina@gmail.com>

* Update manager/api/logging.go

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* Add TTL manager initialization in TestStop

Signed-off-by: Sammy Oina <sammyoina@gmail.com>

* Fix logging format in Shutdown method for consistency

Signed-off-by: Sammy Oina <sammyoina@gmail.com>

* Add unit tests for TTL manager functionality

Signed-off-by: Sammy Oina <sammyoina@gmail.com>

* Enhance TTL tests with mutex for thread safety in expiration checks

Signed-off-by: Sammy Oina <sammyoina@gmail.com>

* Add TTL parameter to CreateVM in TestRun for improved testing scenarios

Signed-off-by: Sammy Oina <sammyoina@gmail.com>

* Add Shutdown test to verify VM cleanup and TTL manager integration

Signed-off-by: Sammy Oina <sammyoina@gmail.com>

---------

Signed-off-by: Sammy Oina <sammyoina@gmail.com>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
2025-07-11 16:21:29 +02:00

343 lines
6.4 KiB
Go

// Copyright (c) Ultraviolet
// SPDX-License-Identifier: Apache-2.0
package manager
import (
"fmt"
"sync"
"testing"
"time"
)
func TestNewTTLManager(t *testing.T) {
tm := NewTTLManager()
if tm == nil {
t.Fatal("NewTTLManager() returned nil")
}
if tm.timers == nil {
t.Fatal("NewTTLManager() did not initialize timers map")
}
if len(tm.timers) != 0 {
t.Errorf("NewTTLManager() timers map should be empty, got %d entries", len(tm.timers))
}
}
func TestSetTTL_Basic(t *testing.T) {
tm := NewTTLManager()
mu := sync.Mutex{}
expired := false
vmID := "test-vm-1"
ttl := 50 * time.Millisecond
cancelFunc := tm.SetTTL(vmID, ttl, func() {
mu.Lock()
defer mu.Unlock()
expired = true
})
tm.mu.RLock()
if _, exists := tm.timers[vmID]; !exists {
t.Error("Timer was not created for VM")
}
tm.mu.RUnlock()
time.Sleep(100 * time.Millisecond)
mu.Lock()
if !expired {
t.Error("TTL did not expire as expected")
}
mu.Unlock()
cancelFunc()
tm.mu.RLock()
if _, exists := tm.timers[vmID]; exists {
t.Error("Timer should be cleaned up after expiry")
}
tm.mu.RUnlock()
}
func TestSetTTL_CancelBeforeExpiry(t *testing.T) {
tm := NewTTLManager()
expired := false
vmID := "test-vm-2"
ttl := 100 * time.Millisecond
cancelFunc := tm.SetTTL(vmID, ttl, func() {
expired = true
})
time.Sleep(20 * time.Millisecond)
cancelFunc()
time.Sleep(150 * time.Millisecond)
if expired {
t.Error("TTL should not have expired after being cancelled")
}
tm.mu.RLock()
if _, exists := tm.timers[vmID]; exists {
t.Error("Timer should be cleaned up after cancellation")
}
tm.mu.RUnlock()
}
func TestSetTTL_OverwriteExistingTimer(t *testing.T) {
tm := NewTTLManager()
mu := sync.Mutex{}
firstExpired := false
secondExpired := false
vmID := "test-vm-3"
// Set first TTL
tm.SetTTL(vmID, 200*time.Millisecond, func() {
firstExpired = true
})
// Immediately overwrite with second TTL
tm.SetTTL(vmID, 50*time.Millisecond, func() {
mu.Lock()
defer mu.Unlock()
secondExpired = true
})
// Wait for second TTL to expire
time.Sleep(100 * time.Millisecond)
if firstExpired {
t.Error("First TTL should not have expired (it was overwritten)")
}
mu.Lock()
if !secondExpired {
t.Error("Second TTL should have expired")
}
mu.Unlock()
// Verify only one timer entry exists (or none after cleanup)
tm.mu.RLock()
count := len(tm.timers)
tm.mu.RUnlock()
if count > 1 {
t.Errorf("Expected at most 1 timer entry, got %d", count)
}
}
func TestSetTTL_MultipleConcurrentTimers(t *testing.T) {
tm := NewTTLManager()
numVMs := 5
expiredCount := int32(0)
var mu sync.Mutex
for i := 0; i < numVMs; i++ {
vmID := fmt.Sprintf("vm-%d", i)
tm.SetTTL(vmID, 50*time.Millisecond, func() {
mu.Lock()
expiredCount++
mu.Unlock()
})
}
tm.mu.RLock()
if len(tm.timers) != numVMs {
t.Errorf("Expected %d timers, got %d", numVMs, len(tm.timers))
}
tm.mu.RUnlock()
time.Sleep(100 * time.Millisecond)
mu.Lock()
finalCount := expiredCount
mu.Unlock()
if int(finalCount) != numVMs {
t.Errorf("Expected %d timers to expire, got %d", numVMs, finalCount)
}
}
func TestCancelTTL_ExistingTimer(t *testing.T) {
tm := NewTTLManager()
expired := false
vmID := "test-vm-4"
tm.SetTTL(vmID, 100*time.Millisecond, func() {
expired = true
})
// Cancel the timer
tm.CancelTTL(vmID)
time.Sleep(150 * time.Millisecond)
if expired {
t.Error("TTL should not have expired after being cancelled")
}
// Verify timer was removed
tm.mu.RLock()
if _, exists := tm.timers[vmID]; exists {
t.Error("Timer should be removed after cancellation")
}
tm.mu.RUnlock()
}
func TestCancelTTL_NonExistentTimer(t *testing.T) {
tm := NewTTLManager()
// Should not panic when cancelling non-existent timer
tm.CancelTTL("non-existent-vm")
// Verify timers map is still empty
tm.mu.RLock()
if len(tm.timers) != 0 {
t.Errorf("Expected empty timers map, got %d entries", len(tm.timers))
}
tm.mu.RUnlock()
}
func TestCancelAll_MultipleTimers(t *testing.T) {
tm := NewTTLManager()
numVMs := 3
expiredCount := int32(0)
var mu sync.Mutex
for i := 0; i < numVMs; i++ {
vmID := fmt.Sprintf("vm-%d", i)
tm.SetTTL(vmID, 200*time.Millisecond, func() {
mu.Lock()
expiredCount++
mu.Unlock()
})
}
tm.mu.RLock()
if len(tm.timers) != numVMs {
t.Errorf("Expected %d timers, got %d", numVMs, len(tm.timers))
}
tm.mu.RUnlock()
tm.CancelAll()
tm.mu.RLock()
if len(tm.timers) != 0 {
t.Errorf("Expected 0 timers after CancelAll, got %d", len(tm.timers))
}
tm.mu.RUnlock()
time.Sleep(250 * time.Millisecond)
mu.Lock()
finalCount := expiredCount
mu.Unlock()
if finalCount != 0 {
t.Errorf("Expected 0 timers to expire after CancelAll, got %d", finalCount)
}
}
func TestCancelAll_EmptyManager(t *testing.T) {
tm := NewTTLManager()
tm.CancelAll()
tm.mu.RLock()
if len(tm.timers) != 0 {
t.Errorf("Expected empty timers map, got %d entries", len(tm.timers))
}
tm.mu.RUnlock()
}
func TestConcurrentAccess(t *testing.T) {
tm := NewTTLManager()
var wg sync.WaitGroup
numGoroutines := 10
for i := 0; i < numGoroutines; i++ {
wg.Add(1)
go func(id int) {
defer wg.Done()
vmID := fmt.Sprintf("concurrent-vm-%d", id)
cancelFunc := tm.SetTTL(vmID, 100*time.Millisecond, func() {})
// Sometimes cancel immediately
if id%2 == 0 {
cancelFunc()
}
}(i)
}
for i := 0; i < numGoroutines; i++ {
wg.Add(1)
go func(id int) {
defer wg.Done()
vmID := fmt.Sprintf("concurrent-vm-%d", id)
time.Sleep(10 * time.Millisecond)
tm.CancelTTL(vmID)
}(i)
}
wg.Wait()
tm.CancelAll()
// This test primarily checks that no race conditions occur
// The actual state at the end is unpredictable due to timing
}
func TestSetTTL_ZeroDuration(t *testing.T) {
tm := NewTTLManager()
mu := sync.Mutex{}
expired := false
vmID := "zero-duration-vm"
tm.SetTTL(vmID, 0, func() {
mu.Lock()
defer mu.Unlock()
expired = true
})
time.Sleep(10 * time.Millisecond)
mu.Lock()
if !expired {
t.Error("TTL with zero duration should expire immediately")
}
mu.Unlock()
}
func TestSetTTL_NegativeDuration(t *testing.T) {
tm := NewTTLManager()
mu := sync.Mutex{}
expired := false
vmID := "negative-duration-vm"
tm.SetTTL(vmID, -100*time.Millisecond, func() {
mu.Lock()
defer mu.Unlock()
expired = true
})
time.Sleep(10 * time.Millisecond)
mu.Lock()
defer mu.Unlock()
if !expired {
t.Error("TTL with negative duration should expire immediately")
}
}