Fix socket leak, clean up single use postgres databases (#2413)

spikecurtis · web-flow · commit a82c0eb56032 · 2022-06-16T09:01:33.000-07:00
* Fix socket leak, clean up single use postgres databases

Signed-off-by: Spike Curtis &lt;spike@coder.com&gt;

* Move migrate close defer until after we know it is not nil

Signed-off-by: Spike Curtis &lt;spike@coder.com&gt;
diff --git a/Makefile b/Makefile
@@ -115,9 +115,9 @@ test: test-clean
 .PHONY: test-postgres
 test-postgres: test-clean
 	DB=ci gotestsum --junitfile="gotests.xml" --packages="./..." -- \
-          -covermode=atomic -coverprofile="gotests.coverage" -timeout=10m \
+          -covermode=atomic -coverprofile="gotests.coverage" -timeout=30m \
           -coverpkg=./...,github.com/coder/coder/codersdk \
-          -count=1 -parallel=1 -race -failfast
+          -count=1 -race -failfast
 
 
 .PHONY: test-postgres-docker
diff --git a/cli/list_test.go b/cli/list_test.go
@@ -5,7 +5,7 @@ import (
 	"testing"
 	"time"
 
-	"github.com/stretchr/testify/require"
+	"github.com/stretchr/testify/assert"
 
 	"github.com/coder/coder/cli/clitest"
 	"github.com/coder/coder/coderd/coderdtest"
@@ -30,13 +30,15 @@ func TestList(t *testing.T) {
 		pty := ptytest.New(t)
 		cmd.SetIn(pty.Input())
 		cmd.SetOut(pty.Output())
-		errC := make(chan error)
+		done := make(chan any)
 		go func() {
-			errC <- cmd.ExecuteContext(ctx)
+			errC := cmd.ExecuteContext(ctx)
+			assert.NoError(t, errC)
+			close(done)
 		}()
 		pty.ExpectMatch(workspace.Name)
 		pty.ExpectMatch("Running")
 		cancelFunc()
-		require.NoError(t, <-errC)
+		<-done
 	})
 }
diff --git a/coderd/database/migrate.go b/coderd/database/migrate.go
@@ -1,6 +1,7 @@
 package database
 
 import (
+	"context"
 	"database/sql"
 	"embed"
 	"errors"
@@ -17,12 +18,21 @@ import (
 var migrations embed.FS
 
 func migrateSetup(db *sql.DB) (source.Driver, *migrate.Migrate, error) {
+	ctx := context.Background()
 	sourceDriver, err := iofs.New(migrations, "migrations")
 	if err != nil {
 		return nil, nil, xerrors.Errorf("create iofs: %w", err)
 	}
 
-	dbDriver, err := postgres.WithInstance(db, &postgres.Config{})
+	// there is a postgres.WithInstance() method that takes the DB instance,
+	// but, when you close the resulting Migrate, it closes the DB, which
+	// we don't want.  Instead, create just a connection that will get closed
+	// when migration is done.
+	conn, err := db.Conn(ctx)
+	if err != nil {
+		return nil, nil, xerrors.Errorf("postgres connection: %w", err)
+	}
+	dbDriver, err := postgres.WithConnection(ctx, conn, &postgres.Config{})
 	if err != nil {
 		return nil, nil, xerrors.Errorf("wrap postgres connection: %w", err)
 	}
@@ -36,11 +46,22 @@ func migrateSetup(db *sql.DB) (source.Driver, *migrate.Migrate, error) {
 }
 
 // MigrateUp runs SQL migrations to ensure the database schema is up-to-date.
-func MigrateUp(db *sql.DB) error {
+func MigrateUp(db *sql.DB) (retErr error) {
 	_, m, err := migrateSetup(db)
 	if err != nil {
 		return xerrors.Errorf("migrate setup: %w", err)
 	}
+	defer func() {
+		srcErr, dbErr := m.Close()
+		if retErr != nil {
+			return
+		}
+		if dbErr != nil {
+			retErr = dbErr
+			return
+		}
+		retErr = srcErr
+	}()
 
 	err = m.Up()
 	if err != nil {
diff --git a/coderd/database/postgres/postgres.go b/coderd/database/postgres/postgres.go
@@ -44,7 +44,13 @@ func Open() (string, func(), error) {
 			return "", nil, xerrors.Errorf("create db: %w", err)
 		}
 
-		return "postgres://postgres:postgres@127.0.0.1:5432/" + dbName + "?sslmode=disable", func() {}, nil
+		deleteDB := func() {
+			ddb, _ := sql.Open("postgres", dbURL)
+			defer ddb.Close()
+			_, _ = ddb.Exec("DROP DATABASE " + dbName)
+		}
+
+		return "postgres://postgres:postgres@127.0.0.1:5432/" + dbName + "?sslmode=disable", deleteDB, nil
 	}
 
 	pool, err := dockertest.NewPool("")

Original file line number	Diff line number	Diff line change
`@@ -44,7 +44,13 @@ func Open() (string, func(), error) {`
`44`	`44`	`return "", nil, xerrors.Errorf("create db: %w", err)`
`45`	`45`	`}`
`46`	`46`
`47`		`- return "postgres://postgres:postgres@127.0.0.1:5432/" + dbName + "?sslmode=disable", func() {}, nil`
	`47`	`+ deleteDB := func() {`
	`48`	`+ ddb, _ := sql.Open("postgres", dbURL)`
	`49`	`+ defer ddb.Close()`
	`50`	`+ _, _ = ddb.Exec("DROP DATABASE " + dbName)`
	`51`	`+ }`
	`52`	`+`
	`53`	`+ return "postgres://postgres:postgres@127.0.0.1:5432/" + dbName + "?sslmode=disable", deleteDB, nil`
`48`	`54`	`}`
`49`	`55`
`50`	`56`	`pool, err := dockertest.NewPool("")`