PMM-12896 Add limit for actions/jobs executed on the same DB at the same time

agent/runner/actions/mongodb_explain_action.go

@@ -18,7 +18,6 @@ import (
 	"context"
 	"fmt"
 	"path/filepath"
-	"strings"
 	"time"
 
 	"github.com/percona/percona-toolkit/src/go/mongolib/proto"
@@ -44,7 +43,7 @@ var errCannotExplain = fmt.Errorf("cannot explain this type of query")
 
 // NewMongoDBExplainAction creates a MongoDB EXPLAIN query Action.
 func NewMongoDBExplainAction(id string, timeout time.Duration, params *agentpb.StartActionRequest_MongoDBExplainParams, tempDir string) (Action, error) {
-	dsn, err := templates.RenderDSN(params.Dsn, params.TextFiles, filepath.Join(tempDir, strings.ToLower(mongoDBExplainActionType), id))
+	dsn, err := templates.RenderDSN(params.Dsn, params.TextFiles, filepath.Join(tempDir, mongoDBExplainActionType, id))
 	if err != nil {
 		return nil, errors.WithStack(err)
 	}

agent/runner/actions/mongodb_query_admincommand_action.go

@@ -17,7 +17,6 @@ package actions
 import (
 	"context"
 	"path/filepath"
-	"strings"
 	"time"
 
 	"github.com/pkg/errors"
@@ -49,7 +48,7 @@ func NewMongoDBQueryAdmincommandAction(
 	arg interface{},
 	tempDir string,
 ) (Action, error) {
-	dsn, err := templates.RenderDSN(dsn, files, filepath.Join(tempDir, strings.ToLower(mongoDBQueryAdminCommandActionType), id))
+	dsn, err := templates.RenderDSN(dsn, files, filepath.Join(tempDir, mongoDBQueryAdminCommandActionType, id))
 	if err != nil {
 		return nil, errors.WithStack(err)
 	}

agent/runner/actions/mysql_explain_action_test.go

@@ -174,7 +174,7 @@ func TestMySQLExplain(t *testing.T) {
 			OutputFormat: agentpb.MysqlExplainOutputFormat_MYSQL_EXPLAIN_OUTPUT_FORMAT_DEFAULT,
 		}
 		a, err := NewMySQLExplainAction("", time.Second, params)
-		assert.Regexp(t, `Query to EXPLAIN is empty`, err.Error())
+		assert.ErrorContains(t, err, `Query to EXPLAIN is empty`)
 		assert.Nil(t, a)
 	})
 
@@ -210,7 +210,7 @@ func TestMySQLExplain(t *testing.T) {
 			OutputFormat: agentpb.MysqlExplainOutputFormat_MYSQL_EXPLAIN_OUTPUT_FORMAT_DEFAULT,
 		}
 		a, err := NewMySQLExplainAction("", time.Second, params)
-		assert.Error(t, err, "EXPLAIN failed because the query was too long and trimmed. Set max-query-length to a larger value.")
+		assert.ErrorContains(t, err, "EXPLAIN failed because the query was too long and trimmed. Set max-query-length to a larger value.")
 		assert.Nil(t, a)
 	})
 

agent/runner/actions/postgresql_query_select_action.go

@@ -50,7 +50,7 @@ func NewPostgreSQLQuerySelectAction(id string, timeout time.Duration, params *ag
 		return nil, errors.New("query contains ';'")
 	}
 
-	dsn, err := templates.RenderDSN(params.Dsn, params.TlsFiles, filepath.Join(tempDir, strings.ToLower(postgreSQLQuerySelectActionType), id))
+	dsn, err := templates.RenderDSN(params.Dsn, params.TlsFiles, filepath.Join(tempDir, postgreSQLQuerySelectActionType, id))
 	if err != nil {
 		return nil, errors.WithStack(err)
 	}

agent/runner/actions/postgresql_query_show_action.go

@@ -18,7 +18,6 @@ import (
 	"context"
 	"database/sql"
 	"path/filepath"
-	"strings"
 	"time"
 
 	"github.com/lib/pq"
@@ -39,7 +38,7 @@ type postgresqlQueryShowAction struct {
 
 // NewPostgreSQLQueryShowAction creates PostgreSQL SHOW query Action.
 func NewPostgreSQLQueryShowAction(id string, timeout time.Duration, params *agentpb.StartActionRequest_PostgreSQLQueryShowParams, tempDir string) (Action, error) {
-	dsn, err := templates.RenderDSN(params.Dsn, params.TlsFiles, filepath.Join(tempDir, strings.ToLower(postgreSQLQueryShowActionType), id))
+	dsn, err := templates.RenderDSN(params.Dsn, params.TlsFiles, filepath.Join(tempDir, postgreSQLQueryShowActionType, id))
 	if err != nil {
 		return nil, errors.WithStack(err)
 	}

agent/runner/actions/postgresql_show_create_table_action.go

@@ -79,7 +79,7 @@ func NewPostgreSQLShowCreateTableAction(
 	params *agentpb.StartActionRequest_PostgreSQLShowCreateTableParams,
 	tempDir string,
 ) (Action, error) {
-	dsn, err := templates.RenderDSN(params.Dsn, params.TlsFiles, filepath.Join(tempDir, strings.ToLower(postgreSQLShowCreateTableActionType), id))
+	dsn, err := templates.RenderDSN(params.Dsn, params.TlsFiles, filepath.Join(tempDir, postgreSQLShowCreateTableActionType, id))
 	if err != nil {
 		return nil, errors.WithStack(err)
 	}

agent/runner/actions/postgresql_show_index_action.go

@@ -42,7 +42,7 @@ type postgresqlShowIndexAction struct {
 // NewPostgreSQLShowIndexAction creates PostgreSQL SHOW INDEX Action.
 // This is an Action that can run `SHOW INDEX` command on PostgreSQL service with given DSN.
 func NewPostgreSQLShowIndexAction(id string, timeout time.Duration, params *agentpb.StartActionRequest_PostgreSQLShowIndexParams, tempDir string) (Action, error) {
-	dsn, err := templates.RenderDSN(params.Dsn, params.TlsFiles, filepath.Join(tempDir, strings.ToLower(postgreSQLShowIndexActionType), id))
+	dsn, err := templates.RenderDSN(params.Dsn, params.TlsFiles, filepath.Join(tempDir, postgreSQLShowIndexActionType, id))
 	if err != nil {
 		return nil, errors.WithStack(err)
 	}

agent/runner/actions/pt_mysql_summary_action.go

@@ -66,6 +66,10 @@ func (a *ptMySQLSummaryAction) Type() string {
 
 // DSN returns a DSN for the Action.
 func (a *ptMySQLSummaryAction) DSN() string {
+	if a.params.Socket != "" {
+		return a.params.Socket
+	}
+
 	return net.JoinHostPort(a.params.Host, strconv.FormatUint(uint64(a.params.Port), 10))
 }
 

agent/runner/runner.go

@@ -127,6 +127,7 @@ func (r *Runner) acquire(ctx context.Context, token string) error {
 	}
 
 	if err := r.gSem.Acquire(ctx, 1); err != nil {
+		r.releaseL(token)
 		return err
 	}
 
@@ -137,6 +138,11 @@ func (r *Runner) acquire(ctx context.Context, token string) error {
 func (r *Runner) release(token string) {
 	r.gSem.Release(1)
 
+	r.releaseL(token)
+}
+
+// releaseL releases local semaphore for given token.
+func (r *Runner) releaseL(token string) {
 	if token != "" {
 		r.lSemsM.Lock()
 
@@ -150,6 +156,13 @@ func (r *Runner) release(token string) {
 	}
 }
 
+// lSemsLen returns number of local semaphores in use.
+func (r *Runner) lSemsLen() int {
+	r.lSemsM.Lock()
+	defer r.lSemsM.Unlock()
+	return len(r.lSems)
+}
+
 // Run starts jobs execution loop. It reads jobs from the channel and starts them in separate goroutines.
 func (r *Runner) Run(ctx context.Context) {
 	for {

agent/runner/runner_test.go

@@ -149,43 +149,43 @@ func TestPerDBInstanceLimit(t *testing.T) {
 	defer cancel()
 	go cr.Run(ctx)
 
-	j1db1 := testJob{id: "test-1", timeout: time.Second, dsn: "postgresql://db1"}
-	j2db1 := testJob{id: "test-2", timeout: time.Second, dsn: "postgresql://db1"}
-	j3db1 := testJob{id: "test-3", timeout: time.Second, dsn: "postgresql://db1"}
-	j1db2 := testJob{id: "test-4", timeout: time.Second, dsn: "postgresql://db2"}
-	j2db2 := testJob{id: "test-5", timeout: time.Second, dsn: "postgresql://db2"}
-	j3db2 := testJob{id: "test-6", timeout: time.Second, dsn: "postgresql://db2"}
+	db1j1 := testJob{id: "test-1", timeout: time.Second, dsn: "postgresql://db1"}
+	db1j2 := testJob{id: "test-2", timeout: time.Second, dsn: "postgresql://db1"}
+	db1j3 := testJob{id: "test-3", timeout: time.Second, dsn: "postgresql://db1"}
+	db2j1 := testJob{id: "test-4", timeout: time.Second, dsn: "postgresql://db2"}
+	db2j2 := testJob{id: "test-5", timeout: time.Second, dsn: "postgresql://db2"}
+	db2j3 := testJob{id: "test-6", timeout: time.Second, dsn: "postgresql://db2"}
 
-	require.NoError(t, cr.StartJob(j1db1))
-	require.NoError(t, cr.StartJob(j1db2))
+	require.NoError(t, cr.StartJob(db1j1))
+	require.NoError(t, cr.StartJob(db2j1))
 
 	// Let jobs to start
 	time.Sleep(200 * time.Millisecond)
 
-	require.NoError(t, cr.StartJob(j2db1))
-	require.NoError(t, cr.StartJob(j2db2))
-	require.NoError(t, cr.StartJob(j3db1))
-	require.NoError(t, cr.StartJob(j3db2))
+	require.NoError(t, cr.StartJob(db1j2))
+	require.NoError(t, cr.StartJob(db2j2))
+	require.NoError(t, cr.StartJob(db1j3))
+	require.NoError(t, cr.StartJob(db2j3))
 
 	// Let rest jobs to reach semaphores
 	time.Sleep(300 * time.Millisecond)
 
-	assert.True(t, cr.IsRunning(j1db1.ID()))
-	assert.True(t, cr.IsRunning(j1db2.ID()))
-	assert.False(t, cr.IsRunning(j2db1.ID()))
-	assert.False(t, cr.IsRunning(j2db2.ID()))
-	assert.False(t, cr.IsRunning(j3db1.ID()))
-	assert.False(t, cr.IsRunning(j3db2.ID()))
+	assert.True(t, cr.IsRunning(db1j1.ID()))
+	assert.True(t, cr.IsRunning(db2j1.ID()))
+	assert.False(t, cr.IsRunning(db1j2.ID()))
+	assert.False(t, cr.IsRunning(db2j2.ID()))
+	assert.False(t, cr.IsRunning(db1j3.ID()))
+	assert.False(t, cr.IsRunning(db2j3.ID()))
 
 	// Over time all jobs are terminated
 	time.Sleep(time.Second)
 
-	assert.False(t, cr.IsRunning(j1db1.ID()))
-	assert.False(t, cr.IsRunning(j1db2.ID()))
-	assert.False(t, cr.IsRunning(j2db1.ID()))
-	assert.False(t, cr.IsRunning(j2db2.ID()))
-	assert.False(t, cr.IsRunning(j3db1.ID()))
-	assert.False(t, cr.IsRunning(j3db2.ID()))
+	assert.False(t, cr.IsRunning(db1j1.ID()))
+	assert.False(t, cr.IsRunning(db2j1.ID()))
+	assert.False(t, cr.IsRunning(db1j2.ID()))
+	assert.False(t, cr.IsRunning(db2j2.ID()))
+	assert.False(t, cr.IsRunning(db1j3.ID()))
+	assert.False(t, cr.IsRunning(db2j3.ID()))
 }
 
 func TestDefaultPerDBInstanceLimit(t *testing.T) {
@@ -196,43 +196,43 @@ func TestDefaultPerDBInstanceLimit(t *testing.T) {
 	defer cancel()
 	go cr.Run(ctx)
 
-	j1db1 := testJob{id: "test-1", timeout: time.Second, dsn: "postgresql://db1"}
-	j2db1 := testJob{id: "test-2", timeout: time.Second, dsn: "postgresql://db1"}
-	j3db1 := testJob{id: "test-3", timeout: time.Second, dsn: "postgresql://db1"}
-	j1db2 := testJob{id: "test-4", timeout: time.Second, dsn: "postgresql://db2"}
-	j2db2 := testJob{id: "test-5", timeout: time.Second, dsn: "postgresql://db2"}
-	j3db2 := testJob{id: "test-6", timeout: time.Second, dsn: "postgresql://db2"}
+	db1j1 := testJob{id: "test-1", timeout: time.Second, dsn: "postgresql://db1"}
+	db1j2 := testJob{id: "test-2", timeout: time.Second, dsn: "postgresql://db1"}
+	db1j3 := testJob{id: "test-3", timeout: time.Second, dsn: "postgresql://db1"}
+	db2j1 := testJob{id: "test-4", timeout: time.Second, dsn: "postgresql://db2"}
+	db2j2 := testJob{id: "test-5", timeout: time.Second, dsn: "postgresql://db2"}
+	db2j3 := testJob{id: "test-6", timeout: time.Second, dsn: "postgresql://db2"}
 
-	require.NoError(t, cr.StartJob(j1db1))
-	require.NoError(t, cr.StartJob(j1db2))
-	require.NoError(t, cr.StartJob(j2db1))
-	require.NoError(t, cr.StartJob(j2db2))
+	require.NoError(t, cr.StartJob(db1j1))
+	require.NoError(t, cr.StartJob(db2j1))
+	require.NoError(t, cr.StartJob(db1j2))
+	require.NoError(t, cr.StartJob(db2j2))
 
 	// Let jobs to start
 	time.Sleep(200 * time.Millisecond)
 
-	require.NoError(t, cr.StartJob(j3db1))
-	require.NoError(t, cr.StartJob(j3db2))
+	require.NoError(t, cr.StartJob(db1j3))
+	require.NoError(t, cr.StartJob(db2j3))
 
 	// Let rest jobs to reach semaphores
 	time.Sleep(300 * time.Millisecond)
 
-	assert.True(t, cr.IsRunning(j1db1.ID()))
-	assert.True(t, cr.IsRunning(j1db2.ID()))
-	assert.True(t, cr.IsRunning(j2db1.ID()))
-	assert.True(t, cr.IsRunning(j2db2.ID()))
-	assert.False(t, cr.IsRunning(j3db1.ID()))
-	assert.False(t, cr.IsRunning(j3db2.ID()))
+	assert.True(t, cr.IsRunning(db1j1.ID()))
+	assert.True(t, cr.IsRunning(db2j1.ID()))
+	assert.True(t, cr.IsRunning(db1j2.ID()))
+	assert.True(t, cr.IsRunning(db2j2.ID()))
+	assert.False(t, cr.IsRunning(db1j3.ID()))
+	assert.False(t, cr.IsRunning(db2j3.ID()))
 
 	// Over time all jobs are terminated
 	time.Sleep(time.Second)
 
-	assert.False(t, cr.IsRunning(j1db1.ID()))
-	assert.False(t, cr.IsRunning(j1db2.ID()))
-	assert.False(t, cr.IsRunning(j2db1.ID()))
-	assert.False(t, cr.IsRunning(j2db2.ID()))
-	assert.False(t, cr.IsRunning(j3db1.ID()))
-	assert.False(t, cr.IsRunning(j3db2.ID()))
+	assert.False(t, cr.IsRunning(db1j1.ID()))
+	assert.False(t, cr.IsRunning(db2j1.ID()))
+	assert.False(t, cr.IsRunning(db1j2.ID()))
+	assert.False(t, cr.IsRunning(db2j2.ID()))
+	assert.False(t, cr.IsRunning(db1j3.ID()))
+	assert.False(t, cr.IsRunning(db2j3.ID()))
 }
 
 func TestConcurrentRunnerTimeout(t *testing.T) {
@@ -329,6 +329,39 @@ func TestConcurrentRunnerCancel(t *testing.T) {
 	assert.Empty(t, cr.cancels)
 }
 
+func TestSemaphoresReleasing(t *testing.T) {
+	t.Parallel()
+	cr := New(1, 1)
+	err := cr.gSem.Acquire(context.TODO(), 1) // Acquire global semaphore to block all jobs
+	require.NoError(t, err)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	go cr.Run(ctx)
+
+	j := testJob{id: "test-1", timeout: time.Second, dsn: "test"}
+
+	require.NoError(t, cr.StartJob(j))
+
+	// Let job to start
+	time.Sleep(200 * time.Millisecond)
+
+	// Check that job is started and local semaphore was acquired
+	assert.Equal(t, cr.lSemsLen(), 1)
+
+	// Check that job is not running, because it's waiting for global semaphore to be acquired
+	assert.False(t, cr.IsRunning(j.ID()))
+
+	// Cancel context to stop job
+	cancel()
+
+	// Let job to start and release resources
+	time.Sleep(200 * time.Millisecond)
+
+	// Check that local samaphore was released
+	assert.Equal(t, cr.lSemsLen(), 0)
+}
+
 type testJob struct {
 	id      string
 	timeout time.Duration