Check for pool max runners in CreateInstance tx

This change moves the check for max runners within the CreateInstance
function, which will check that the pool max runners is not yet reached
within a transaction before creating a new instance.

Signed-off-by: Gabriel Adrian Samfira <gsamfira@cloudbasesolutions.com>

Author: gabriel-samfira

config/config.go

@@ -569,7 +569,7 @@ func (s *SQLite) Validate() error {
 }
 
 func (s *SQLite) ConnectionString() (string, error) {
-	connectionString := fmt.Sprintf("%s?_journal_mode=WAL&_foreign_keys=ON", s.DBFile)
+	connectionString := fmt.Sprintf("%s?_journal_mode=WAL&_foreign_keys=ON&_txlock=immediate", s.DBFile)
 	if s.BusyTimeoutSeconds > 0 {
 		timeout := s.BusyTimeoutSeconds * 1000
 		connectionString = fmt.Sprintf("%s&_busy_timeout=%d", connectionString, timeout)

config/config_test.go

@@ -387,13 +387,13 @@ func TestGormParams(t *testing.T) {
 	dbType, uri, err := cfg.GormParams()
 	require.Nil(t, err)
 	require.Equal(t, SQLiteBackend, dbType)
-	require.Equal(t, filepath.Join(dir, "garm.db?_journal_mode=WAL&_foreign_keys=ON"), uri)
+	require.Equal(t, filepath.Join(dir, "garm.db?_journal_mode=WAL&_foreign_keys=ON&_txlock=immediate"), uri)
 
 	cfg.SQLite.BusyTimeoutSeconds = 5
 	dbType, uri, err = cfg.GormParams()
 	require.Nil(t, err)
 	require.Equal(t, SQLiteBackend, dbType)
-	require.Equal(t, filepath.Join(dir, "garm.db?_journal_mode=WAL&_foreign_keys=ON&_busy_timeout=5000"), uri)
+	require.Equal(t, filepath.Join(dir, "garm.db?_journal_mode=WAL&_foreign_keys=ON&_txlock=immediate&_busy_timeout=5000"), uri)
 
 	cfg.DbBackend = MySQLBackend
 	cfg.MySQL = getMySQLDefaultConfig()

database/sql/instances.go

@@ -17,7 +17,9 @@ package sql
 import (
 	"context"
 	"encoding/json"
+	"fmt"
 	"log/slog"
+	"math"
 
 	"github.com/google/uuid"
 	"github.com/pkg/errors"
@@ -30,57 +32,76 @@ import (
 	"github.com/cloudbase/garm/params"
 )
 
-func (s *sqlDatabase) CreateInstance(_ context.Context, poolID string, param params.CreateInstanceParams) (instance params.Instance, err error) {
+func (s *sqlDatabase) CreateInstance(ctx context.Context, poolID string, param params.CreateInstanceParams) (instance params.Instance, err error) {
 	s.writeMux.Lock()
 	defer s.writeMux.Unlock()
 
-	pool, err := s.getPoolByID(s.conn, poolID)
-	if err != nil {
-		return params.Instance{}, errors.Wrap(err, "fetching pool")
-	}
-
 	defer func() {
 		if err == nil {
 			s.sendNotify(common.InstanceEntityType, common.CreateOperation, instance)
 		}
 	}()
 
-	var labels datatypes.JSON
-	if len(param.AditionalLabels) > 0 {
-		labels, err = json.Marshal(param.AditionalLabels)
+	err = s.conn.Transaction(func(tx *gorm.DB) error {
+		pool, err := s.getPoolByID(tx, poolID)
 		if err != nil {
-			return params.Instance{}, errors.Wrap(err, "marshalling labels")
+			return errors.Wrap(err, "fetching pool")
+		}
+		var cnt int64
+		q := tx.Model(&Instance{}).Where("pool_id = ?", pool.ID).Count(&cnt)
+		if q.Error != nil {
+			return fmt.Errorf("error fetching instance count: %w", q.Error)
 		}
-	}
 
-	var secret []byte
-	if len(param.JitConfiguration) > 0 {
-		secret, err = s.marshalAndSeal(param.JitConfiguration)
-		if err != nil {
-			return params.Instance{}, errors.Wrap(err, "marshalling jit config")
+		var maxRunners int64
+		if pool.MaxRunners > math.MaxInt64 {
+			maxRunners = math.MaxInt64
+		} else {
+			maxRunners = int64(pool.MaxRunners)
+		}
+		if cnt >= maxRunners {
+			return runnerErrors.NewConflictError("max runners reached for pool %s", pool.ID)
+		}
+		var labels datatypes.JSON
+		if len(param.AditionalLabels) > 0 {
+			labels, err = json.Marshal(param.AditionalLabels)
+			if err != nil {
+				return errors.Wrap(err, "marshalling labels")
+			}
 		}
-	}
 
-	newInstance := Instance{
-		Pool:              pool,
-		Name:              param.Name,
-		Status:            param.Status,
-		RunnerStatus:      param.RunnerStatus,
-		OSType:            param.OSType,
-		OSArch:            param.OSArch,
-		CallbackURL:       param.CallbackURL,
-		MetadataURL:       param.MetadataURL,
-		GitHubRunnerGroup: param.GitHubRunnerGroup,
-		JitConfiguration:  secret,
-		AditionalLabels:   labels,
-		AgentID:           param.AgentID,
-	}
-	q := s.conn.Create(&newInstance)
-	if q.Error != nil {
-		return params.Instance{}, errors.Wrap(q.Error, "creating instance")
+		var secret []byte
+		if len(param.JitConfiguration) > 0 {
+			secret, err = s.marshalAndSeal(param.JitConfiguration)
+			if err != nil {
+				return errors.Wrap(err, "marshalling jit config")
+			}
+		}
+		newInstance := Instance{
+			Pool:              pool,
+			Name:              param.Name,
+			Status:            param.Status,
+			RunnerStatus:      param.RunnerStatus,
+			OSType:            param.OSType,
+			OSArch:            param.OSArch,
+			CallbackURL:       param.CallbackURL,
+			MetadataURL:       param.MetadataURL,
+			GitHubRunnerGroup: param.GitHubRunnerGroup,
+			JitConfiguration:  secret,
+			AditionalLabels:   labels,
+			AgentID:           param.AgentID,
+		}
+		q = tx.Create(&newInstance)
+		if q.Error != nil {
+			return errors.Wrap(q.Error, "creating instance")
+		}
+		return nil
+	})
+	if err != nil {
+		return params.Instance{}, errors.Wrap(err, "creating instance")
 	}
 
-	return s.sqlToParamsInstance(newInstance)
+	return s.GetInstanceByName(ctx, param.Name)
 }
 
 func (s *sqlDatabase) getPoolInstanceByName(poolID string, instanceName string) (Instance, error) {

database/sql/instances_test.go

@@ -204,17 +204,48 @@ func (s *InstancesTestSuite) TestCreateInstance() {
 func (s *InstancesTestSuite) TestCreateInstanceInvalidPoolID() {
 	_, err := s.Store.CreateInstance(s.adminCtx, "dummy-pool-id", params.CreateInstanceParams{})
 
-	s.Require().Equal("fetching pool: parsing id: invalid request", err.Error())
+	s.Require().Equal("creating instance: fetching pool: parsing id: invalid request", err.Error())
+}
+
+func (s *InstancesTestSuite) TestCreateInstanceMaxRunnersReached() {
+	// Pool has MaxRunners=4 and already has 3 instances
+	// Create one more to reach the limit
+	_, err := s.Store.CreateInstance(s.adminCtx, s.Fixtures.Pool.ID, params.CreateInstanceParams{
+		Name:         "test-instance-4",
+		OSType:       "linux",
+		OSArch:       "amd64",
+		CallbackURL:  "https://garm.example.com/",
+		Status:       commonParams.InstanceRunning,
+		RunnerStatus: params.RunnerIdle,
+	})
+	s.Require().Nil(err)
+
+	// Now try to create a 5th instance, which should fail
+	_, err = s.Store.CreateInstance(s.adminCtx, s.Fixtures.Pool.ID, params.CreateInstanceParams{
+		Name:         "test-instance-5",
+		OSType:       "linux",
+		OSArch:       "amd64",
+		CallbackURL:  "https://garm.example.com/",
+		Status:       commonParams.InstanceRunning,
+		RunnerStatus: params.RunnerIdle,
+	})
+
+	s.Require().NotNil(err)
+	s.Require().Contains(err.Error(), "max runners reached")
 }
 
 func (s *InstancesTestSuite) TestCreateInstanceDBCreateErr() {
 	pool := s.Fixtures.Pool
 
+	s.Fixtures.SQLMock.ExpectBegin()
 	s.Fixtures.SQLMock.
 		ExpectQuery(regexp.QuoteMeta("SELECT * FROM `pools` WHERE id = ? AND `pools`.`deleted_at` IS NULL ORDER BY `pools`.`id` LIMIT ?")).
 		WithArgs(pool.ID, 1).
-		WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(pool.ID))
-	s.Fixtures.SQLMock.ExpectBegin()
+		WillReturnRows(sqlmock.NewRows([]string{"id", "max_runners"}).AddRow(pool.ID, pool.MaxRunners))
+	s.Fixtures.SQLMock.
+		ExpectQuery(regexp.QuoteMeta("SELECT count(*) FROM `instances` WHERE pool_id = ? AND `instances`.`deleted_at` IS NULL")).
+		WithArgs(pool.ID).
+		WillReturnRows(sqlmock.NewRows([]string{"count(*)"}).AddRow(0))
 	s.Fixtures.SQLMock.
 		ExpectExec("INSERT INTO `pools`").
 		WillReturnResult(sqlmock.NewResult(1, 1))
@@ -227,7 +258,7 @@ func (s *InstancesTestSuite) TestCreateInstanceDBCreateErr() {
 
 	s.assertSQLMockExpectations()
 	s.Require().NotNil(err)
-	s.Require().Equal("creating instance: mocked insert instance error", err.Error())
+	s.Require().Equal("creating instance: creating instance: mocked insert instance error", err.Error())
 }
 
 func (s *InstancesTestSuite) TestGetPoolInstanceByName() {

database/watcher/watcher_store_test.go

@@ -167,6 +167,7 @@ func (s *WatcherStoreTestSuite) TestInstanceWatcher() {
 		ProviderName: "test-provider",
 		Image:        "test-image",
 		Flavor:       "test-flavor",
+		MaxRunners:   100,
 		OSType:       commonParams.Linux,
 		OSArch:       commonParams.Amd64,
 		Tags:         []string{"test-tag"},

runner/pool/pool.go

@@ -1118,15 +1118,6 @@ func (r *basePoolManager) addRunnerToPool(pool params.Pool, aditionalLabels []st
 		return fmt.Errorf("pool %s is disabled", pool.ID)
 	}
 
-	poolInstanceCount, err := r.store.PoolInstanceCount(r.ctx, pool.ID)
-	if err != nil {
-		return fmt.Errorf("failed to list pool instances: %w", err)
-	}
-
-	if poolInstanceCount >= int64(pool.MaxRunners) {
-		return fmt.Errorf("max workers (%d) reached for pool %s", pool.MaxRunners, pool.ID)
-	}
-
 	if err := r.AddRunner(r.ctx, pool.ID, aditionalLabels); err != nil {
 		return fmt.Errorf("failed to add new instance for pool %s: %s", pool.ID, err)
 	}