Fix AA instance heartbeat deduplication

Enforce the AA heartbeat instance_id uniqueness migration and clean up
historical duplicate heartbeat rows before creating the index. Make active
instance listing resilient to stale duplicates, and update heartbeat upserts
to refresh the full instance metadata including IP and EAS fields.

Signed-off-by: Jiale Zhang <xinjian.zjl@alibaba-inc.com>

Author: Jiale Zhang

trustee-gateway/internal/persistence/repository/aa_repository.go

@@ -33,6 +33,10 @@ func (r *AAInstanceRepository) UpsertHeartbeat(heartbeat *models.AAInstanceHeart
 			"image_id",
 			"instance_name",
 			"owner_account_id",
+			"eas_model_id",
+			"eas_instance_id",
+			"eas_pod_name",
+			"ip",
 			"client_ip",
 			"last_heartbeat",
 			"updated_at",
@@ -43,8 +47,22 @@ func (r *AAInstanceRepository) UpsertHeartbeat(heartbeat *models.AAInstanceHeart
 // GetActiveHeartbeats retrieves all heartbeats that are newer than the cutoff time
 func (r *AAInstanceRepository) GetActiveHeartbeats(cutoffTime time.Time) ([]models.AAInstanceHeartbeat, error) {
 	var heartbeats []models.AAInstanceHeartbeat
-	err := r.db.Where("last_heartbeat >= ?", cutoffTime).Order("last_heartbeat desc").Find(&heartbeats).Error
-	return heartbeats, err
+	err := r.db.Where("last_heartbeat >= ?", cutoffTime).Order("last_heartbeat desc, id desc").Find(&heartbeats).Error
+	if err != nil {
+		return nil, err
+	}
+
+	seen := make(map[string]struct{}, len(heartbeats))
+	deduped := make([]models.AAInstanceHeartbeat, 0, len(heartbeats))
+	for _, heartbeat := range heartbeats {
+		if _, ok := seen[heartbeat.InstanceID]; ok {
+			continue
+		}
+		seen[heartbeat.InstanceID] = struct{}{}
+		deduped = append(deduped, heartbeat)
+	}
+
+	return deduped, nil
 }
 
 // CleanupExpiredHeartbeats removes heartbeat records older than the cutoff time

trustee-gateway/internal/persistence/repository/aa_repository_test.go

@@ -0,0 +1,109 @@
+package repository
+
+import (
+	"testing"
+	"time"
+
+	"github.com/openanolis/trustee/gateway/internal/models"
+	"github.com/openanolis/trustee/gateway/internal/persistence/storage"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"gorm.io/driver/sqlite"
+	"gorm.io/gorm"
+)
+
+func setupAAInstanceTestDB(t *testing.T, withUniqueIndex bool) *storage.Database {
+	t.Helper()
+
+	db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
+	require.NoError(t, err)
+
+	err = db.AutoMigrate(&models.AAInstanceHeartbeat{})
+	require.NoError(t, err)
+
+	if withUniqueIndex {
+		err = db.Exec("CREATE UNIQUE INDEX idx_aa_heartbeat_instance_id ON aa_instance_heartbeats(instance_id)").Error
+		require.NoError(t, err)
+	}
+
+	return &storage.Database{DB: db}
+}
+
+func TestAAInstanceRepositoryUpsertHeartbeatUpdatesExistingRecord(t *testing.T) {
+	repo := NewAAInstanceRepository(setupAAInstanceTestDB(t, true))
+
+	err := repo.UpsertHeartbeat(&models.AAInstanceHeartbeat{
+		InstanceInfo: models.InstanceInfo{
+			InstanceID:     "i-test",
+			ImageID:        "old-image",
+			InstanceName:   "old-name",
+			OwnerAccountID: "old-owner",
+			IP:             "10.0.0.1",
+		},
+		ClientIP: "10.0.0.1",
+	})
+	require.NoError(t, err)
+
+	err = repo.UpsertHeartbeat(&models.AAInstanceHeartbeat{
+		InstanceInfo: models.InstanceInfo{
+			InstanceID:     "i-test",
+			ImageID:        "new-image",
+			InstanceName:   "new-name",
+			OwnerAccountID: "new-owner",
+			EasModelID:     "model-1",
+			EasInstanceID:  "eas-1",
+			EasPodName:     "pod-1",
+			IP:             "10.0.0.2",
+		},
+		ClientIP: "10.0.0.2",
+	})
+	require.NoError(t, err)
+
+	var heartbeats []models.AAInstanceHeartbeat
+	err = repo.db.Find(&heartbeats).Error
+	require.NoError(t, err)
+	require.Len(t, heartbeats, 1)
+
+	heartbeat := heartbeats[0]
+	assert.Equal(t, "i-test", heartbeat.InstanceID)
+	assert.Equal(t, "new-image", heartbeat.ImageID)
+	assert.Equal(t, "new-name", heartbeat.InstanceName)
+	assert.Equal(t, "new-owner", heartbeat.OwnerAccountID)
+	assert.Equal(t, "model-1", heartbeat.EasModelID)
+	assert.Equal(t, "eas-1", heartbeat.EasInstanceID)
+	assert.Equal(t, "pod-1", heartbeat.EasPodName)
+	assert.Equal(t, "10.0.0.2", heartbeat.IP)
+	assert.Equal(t, "10.0.0.2", heartbeat.ClientIP)
+}
+
+func TestAAInstanceRepositoryGetActiveHeartbeatsDeduplicatesInstanceID(t *testing.T) {
+	repo := NewAAInstanceRepository(setupAAInstanceTestDB(t, false))
+	now := time.Now()
+
+	records := []models.AAInstanceHeartbeat{
+		{
+			InstanceInfo:  models.InstanceInfo{InstanceID: "i-test", IP: "10.0.0.1"},
+			ClientIP:      "10.0.0.1",
+			LastHeartbeat: now.Add(-time.Minute),
+		},
+		{
+			InstanceInfo:  models.InstanceInfo{InstanceID: "i-test", IP: "10.0.0.2"},
+			ClientIP:      "10.0.0.2",
+			LastHeartbeat: now,
+		},
+		{
+			InstanceInfo:  models.InstanceInfo{InstanceID: "i-other", IP: "10.0.0.3"},
+			ClientIP:      "10.0.0.3",
+			LastHeartbeat: now.Add(-30 * time.Second),
+		},
+	}
+	require.NoError(t, repo.db.Create(&records).Error)
+
+	heartbeats, err := repo.GetActiveHeartbeats(now.Add(-10 * time.Minute))
+	require.NoError(t, err)
+	require.Len(t, heartbeats, 2)
+
+	assert.Equal(t, "i-test", heartbeats[0].InstanceID)
+	assert.Equal(t, "10.0.0.2", heartbeats[0].ClientIP)
+	assert.Equal(t, "i-other", heartbeats[1].InstanceID)
+}

trustee-gateway/internal/persistence/storage/database.go

@@ -167,12 +167,14 @@ func (d *Database) migrateSchema() error {
 	// This prevents duplicate heartbeat records for the same instance
 	if d.config.Type == "mysql" {
 		if err := d.migrateAAInstanceHeartbeatIndexMySQL(); err != nil {
-			logrus.Warnf("Failed to migrate AAInstanceHeartbeat index: %v", err)
+			return fmt.Errorf("failed to migrate AAInstanceHeartbeat index: %w", err)
 		}
 	} else {
 		// SQLite - file-based SQLite has built-in file locking
 		// In-memory SQLite is per-process, so no cross-process conflict
-		d.migrateAAInstanceHeartbeatIndexSQLite()
+		if err := d.migrateAAInstanceHeartbeatIndexSQLite(); err != nil {
+			return fmt.Errorf("failed to migrate AAInstanceHeartbeat index: %w", err)
+		}
 	}
 
 	return nil
@@ -190,8 +192,7 @@ func (d *Database) migrateAAInstanceHeartbeatIndexMySQL() error {
 	}
 
 	if lockResult != 1 {
-		logrus.Info("Another instance is running migration, skipping")
-		return nil
+		return fmt.Errorf("failed to acquire migration lock: another instance is running migration")
 	}
 
 	// Ensure lock is released when done
@@ -213,18 +214,26 @@ func (d *Database) migrateAAInstanceHeartbeatIndexMySQL() error {
 
 	// Alter column to VARCHAR(255) for index compatibility
 	if err := d.DB.Exec("ALTER TABLE aa_instance_heartbeats MODIFY instance_id VARCHAR(255)").Error; err != nil {
-		logrus.Warnf("Failed to alter instance_id column: %v", err)
+		return fmt.Errorf("failed to alter instance_id column: %w", err)
+	}
+
+	if err := d.DB.Exec("DELETE FROM aa_instance_heartbeats WHERE instance_id IS NULL OR instance_id = ''").Error; err != nil {
+		return fmt.Errorf("failed to clean up empty instance_id records: %w", err)
 	}
 
 	// Clean up duplicate instance_id records before creating unique index
-	// Keep only the latest record (by id) for each instance_id
+	// Keep only the latest heartbeat for each instance_id.
 	cleanupSQL := `
 		DELETE t1 FROM aa_instance_heartbeats t1
 		INNER JOIN aa_instance_heartbeats t2
-		WHERE t1.instance_id = t2.instance_id AND t1.id < t2.id
+		WHERE t1.instance_id = t2.instance_id
+		  AND (
+		    t1.last_heartbeat < t2.last_heartbeat
+		    OR (t1.last_heartbeat = t2.last_heartbeat AND t1.id < t2.id)
+		  )
 	`
 	if err := d.DB.Exec(cleanupSQL).Error; err != nil {
-		logrus.Warnf("Failed to clean up duplicate instance_id records: %v", err)
+		return fmt.Errorf("failed to clean up duplicate instance_id records: %w", err)
 	} else {
 		logrus.Info("Cleaned up duplicate instance_id records in aa_instance_heartbeats")
 	}
@@ -238,29 +247,45 @@ func (d *Database) migrateAAInstanceHeartbeatIndexMySQL() error {
 }
 
 // migrateAAInstanceHeartbeatIndexSQLite handles unique index creation for SQLite
-func (d *Database) migrateAAInstanceHeartbeatIndexSQLite() {
+func (d *Database) migrateAAInstanceHeartbeatIndexSQLite() error {
 	migrator := d.DB.Migrator()
 	if migrator.HasIndex(&models.AAInstanceHeartbeat{}, "idx_aa_heartbeat_instance_id") {
-		return
+		return nil
+	}
+
+	if err := d.DB.Exec("DELETE FROM aa_instance_heartbeats WHERE instance_id IS NULL OR instance_id = ''").Error; err != nil {
+		return fmt.Errorf("failed to clean up empty instance_id records: %w", err)
 	}
 
 	// Clean up duplicate instance_id records before creating unique index
-	// Keep only the latest record (by id) for each instance_id
+	// Keep only the latest heartbeat for each instance_id.
 	cleanupSQL := `
 		DELETE FROM aa_instance_heartbeats
 		WHERE id NOT IN (
-			SELECT MAX(id) FROM aa_instance_heartbeats GROUP BY instance_id
+			SELECT id FROM (
+				SELECT h1.id
+				FROM aa_instance_heartbeats h1
+				LEFT JOIN aa_instance_heartbeats h2
+				  ON h1.instance_id = h2.instance_id
+				 AND (
+				   h1.last_heartbeat < h2.last_heartbeat
+				   OR (h1.last_heartbeat = h2.last_heartbeat AND h1.id < h2.id)
+				 )
+				WHERE h2.id IS NULL
+			)
 		)
 	`
 	if err := d.DB.Exec(cleanupSQL).Error; err != nil {
-		logrus.Warnf("Failed to clean up duplicate instance_id records: %v", err)
+		return fmt.Errorf("failed to clean up duplicate instance_id records: %w", err)
 	} else {
 		logrus.Info("Cleaned up duplicate instance_id records in aa_instance_heartbeats")
 	}
 
 	if err := d.DB.Exec("CREATE UNIQUE INDEX IF NOT EXISTS idx_aa_heartbeat_instance_id ON aa_instance_heartbeats(instance_id)").Error; err != nil {
-		logrus.Warnf("Failed to create unique index on instance_id: %v", err)
+		return fmt.Errorf("failed to create unique index on instance_id: %w", err)
 	}
+
+	return nil
 }
 
 // restoreFromBackup restores data from backup file to in-memory database using native backup API

trustee-gateway/internal/persistence/storage/database_test.go

@@ -0,0 +1,65 @@
+package storage
+
+import (
+	"testing"
+	"time"
+
+	"github.com/openanolis/trustee/gateway/internal/config"
+	"github.com/openanolis/trustee/gateway/internal/models"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"gorm.io/driver/sqlite"
+	"gorm.io/gorm"
+)
+
+func TestMigrateAAInstanceHeartbeatIndexSQLiteCleansDuplicates(t *testing.T) {
+	db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
+	require.NoError(t, err)
+
+	require.NoError(t, db.AutoMigrate(&models.AAInstanceHeartbeat{}))
+
+	now := time.Now()
+	records := []models.AAInstanceHeartbeat{
+		{
+			InstanceInfo:  models.InstanceInfo{InstanceID: "i-test", IP: "10.0.0.1"},
+			ClientIP:      "10.0.0.1",
+			LastHeartbeat: now.Add(-time.Minute),
+		},
+		{
+			InstanceInfo:  models.InstanceInfo{InstanceID: "i-test", IP: "10.0.0.2"},
+			ClientIP:      "10.0.0.2",
+			LastHeartbeat: now,
+		},
+		{
+			InstanceInfo:  models.InstanceInfo{InstanceID: ""},
+			ClientIP:      "10.0.0.3",
+			LastHeartbeat: now,
+		},
+	}
+	require.NoError(t, db.Create(&records).Error)
+
+	database := &Database{
+		DB:     db,
+		config: &config.DatabaseConfig{Type: "sqlite"},
+	}
+	require.NoError(t, database.migrateAAInstanceHeartbeatIndexSQLite())
+
+	var heartbeats []models.AAInstanceHeartbeat
+	require.NoError(t, db.Order("last_heartbeat desc").Find(&heartbeats).Error)
+	require.Len(t, heartbeats, 1)
+	assert.Equal(t, "i-test", heartbeats[0].InstanceID)
+	assert.Equal(t, "10.0.0.2", heartbeats[0].ClientIP)
+
+	require.NoError(t, db.Create(&models.AAInstanceHeartbeat{
+		InstanceInfo:  models.InstanceInfo{InstanceID: "i-other", IP: "10.0.0.4"},
+		ClientIP:      "10.0.0.4",
+		LastHeartbeat: now,
+	}).Error)
+
+	err = db.Create(&models.AAInstanceHeartbeat{
+		InstanceInfo:  models.InstanceInfo{InstanceID: "i-test", IP: "10.0.0.5"},
+		ClientIP:      "10.0.0.5",
+		LastHeartbeat: now,
+	}).Error
+	assert.Error(t, err)
+}