fix: optimize active session count query

This commit adds an index on `tenant_id` in the `active_sessions`
collection to improve the performance of active session counting.
It also updates code in MongoDB store to utilize this index.

Closes shellhub-io#3228

api/store/mongo/migrations/main.go

@@ -69,6 +69,7 @@ func GenerateMigrations() []migrate.Migration {
 		migration57,
 		migration58,
 		migration59,
+		migration60,
 	}
 }
 

api/store/mongo/migrations/migration_60.go

@@ -0,0 +1,48 @@
+package migrations
+
+import (
+	"context"
+
+	"github.com/sirupsen/logrus"
+	migrate "github.com/xakep666/mongo-migrate"
+	"go.mongodb.org/mongo-driver/bson"
+	"go.mongodb.org/mongo-driver/mongo"
+	"go.mongodb.org/mongo-driver/mongo/options"
+)
+
+var migration60 = migrate.Migration{
+	Version:     60,
+	Description: "create index for tenant_id on active_sessions",
+	Up: func(database *mongo.Database) error {
+		logrus.WithFields(logrus.Fields{
+			"component": "migration",
+			"version":   60,
+			"action":    "Up",
+		}).Info("Applying migration up")
+		indexName := "tenant_id"
+		if _, err := database.Collection("active_sessions").Indexes().CreateOne(context.Background(), mongo.IndexModel{
+			Keys: bson.M{
+				"tenant_id": 1,
+			},
+			Options: &options.IndexOptions{ //nolint:exhaustruct
+				Name: &indexName,
+			},
+		}); err != nil {
+			return err
+		}
+
+		return nil
+	},
+	Down: func(database *mongo.Database) error {
+		logrus.WithFields(logrus.Fields{
+			"component": "migration",
+			"version":   60,
+			"action":    "Down",
+		}).Info("Applying migration down")
+		if _, err := database.Collection("active_sessions").Indexes().DropOne(context.Background(), "tenant_id"); err != nil {
+			return err
+		}
+
+		return nil
+	},
+}

api/store/mongo/migrations/migration_60_test.go

@@ -0,0 +1,109 @@
+package migrations
+
+import (
+	"context"
+	"errors"
+	"testing"
+
+	"github.com/shellhub-io/shellhub/api/pkg/dbtest"
+	"github.com/shellhub-io/shellhub/pkg/envs"
+	envMocks "github.com/shellhub-io/shellhub/pkg/envs/mocks"
+	"github.com/sirupsen/logrus"
+	"github.com/stretchr/testify/assert"
+	migrate "github.com/xakep666/mongo-migrate"
+	"go.mongodb.org/mongo-driver/bson"
+)
+
+func TestMigration60(t *testing.T) {
+	logrus.Info("Testing Migration 60")
+
+	db := dbtest.DBServer{}
+	defer db.Stop()
+
+	mock := &envMocks.Backend{}
+	envs.DefaultBackend = mock
+
+	cases := []struct {
+		description string
+		test        func() error
+	}{
+		{
+			"Success to apply up on migration 60",
+			func() error {
+				mock.On("Get", "SHELLHUB_CLOUD").Return("true").Once()
+
+				migrations := GenerateMigrations()[59:60]
+				migrates := migrate.NewMigrate(db.Client().Database("test"), migrations...)
+				err := migrates.Up(migrate.AllAvailable)
+				if err != nil {
+					return err
+				}
+
+				cursor, err := db.Client().Database("test").Collection("active_sessions").Indexes().List(context.Background())
+				if err != nil {
+					return err
+				}
+
+				var found bool
+				for cursor.Next(context.Background()) {
+					var index bson.M
+					if err := cursor.Decode(&index); err != nil {
+						return err
+					}
+
+					if index["name"] == "tenant_id" {
+						found = true
+					}
+				}
+
+				if !found {
+					return errors.New("index not created")
+				}
+
+				return nil
+			},
+		},
+		{
+			"Success to apply down on migration 60",
+			func() error {
+				migrations := GenerateMigrations()[59:60]
+				migrates := migrate.NewMigrate(db.Client().Database("test"), migrations...)
+				err := migrates.Down(migrate.AllAvailable)
+				if err != nil {
+					return err
+				}
+
+				cursor, err := db.Client().Database("test").Collection("active_sessions").Indexes().List(context.Background())
+				if err != nil {
+					return errors.New("index not dropped")
+				}
+
+				var found bool
+				for cursor.Next(context.Background()) {
+					var index bson.M
+					if err := cursor.Decode(&index); err != nil {
+						return err
+					}
+
+					if index["name"] == "tenant_id" {
+						found = true
+					}
+				}
+
+				if found {
+					return errors.New("index not dropped")
+				}
+
+				return nil
+			},
+		},
+	}
+
+	for _, test := range cases {
+		tc := test
+		t.Run(tc.description, func(t *testing.T) {
+			err := tc.test()
+			assert.NoError(t, err)
+		})
+	}
+}

api/store/mongo/session_store.go

@@ -186,6 +186,7 @@ func (s *Store) SessionCreate(ctx context.Context, session models.Session) (*mod
 	as := &models.ActiveSession{
 		UID:      models.UID(session.UID),
 		LastSeen: session.StartedAt,
+		TenantID: session.TenantID,
 	}
 
 	if _, err := s.db.Collection("active_sessions").InsertOne(ctx, &as); err != nil {
@@ -215,12 +216,7 @@ func (s *Store) SessionSetLastSeen(ctx context.Context, uid models.UID) error {
 		return FromMongoError(err)
 	}
 
-	activeSession := &models.ActiveSession{
-		UID:      uid,
-		LastSeen: clock.Now(),
-	}
-
-	if _, err := s.db.Collection("active_sessions").InsertOne(ctx, &activeSession); err != nil {
+	if _, err := s.db.Collection("active_sessions").UpdateOne(ctx, bson.M{"uid": uid}, bson.M{"$set": bson.M{"last_seen": clock.Now()}}); err != nil {
 		return FromMongoError(err)
 	}
 

api/store/mongo/stats_store.go

@@ -105,26 +105,7 @@ func (s *Store) GetStats(ctx context.Context) (*models.Stats, error) {
 		return nil, err
 	}
 
-	query = []bson.M{
-		{
-			"$lookup": bson.M{
-				"from":         "active_sessions",
-				"localField":   "uid",
-				"foreignField": "uid",
-				"as":           "active",
-			},
-		},
-		{
-			"$addFields": bson.M{
-				"active": bson.M{"$anyElementTrue": []interface{}{"$active"}},
-			},
-		},
-		{
-			"$match": bson.M{
-				"active": true,
-			},
-		},
-	}
+	query = []bson.M{}
 
 	// Only match for the respective tenant if requested
 	if tenant := gateway.TenantFromContext(ctx); tenant != nil {
@@ -139,7 +120,7 @@ func (s *Store) GetStats(ctx context.Context) (*models.Stats, error) {
 		"$count": "count",
 	})
 
-	activeSessions, err := AggregateCount(ctx, s.db.Collection("sessions"), query)
+	activeSessions, err := AggregateCount(ctx, s.db.Collection("active_sessions"), query)
 	if err != nil {
 		return nil, err
 	}

pkg/models/session.go

@@ -30,6 +30,7 @@ type Session struct {
 type ActiveSession struct {
 	UID      UID       `json:"uid"`
 	LastSeen time.Time `json:"last_seen" bson:"last_seen"`
+	TenantID string    `json:"tenant_id" bson:"tenant_id"`
 }
 
 type RecordedSession struct {