clean up prepared transactions in (*DB).Close

Taskfile.yml

@@ -238,7 +238,7 @@ tasks:
     cmds:
       - go test ./core/... -tags=ext_test -count=1 -race
       - CGO_ENABLED=1 go test ./parse/... -tags=ext_test -count=1 -race
-      - CGO_ENABLED=1 go test ./... -tags=ext_test -count=1 -race
+      - CGO_ENABLED=1 go test ./... -tags=ext_test,pglive -count=1 -race
 
   # test:it:
   #   desc: Run integration tests ('short' mode)

app/node/build.go

@@ -152,6 +152,8 @@ func (mt *mysteryThing) Price(ctx context.Context, db sql.DB, tx *types.Transact
 }
 
 func buildDB(ctx context.Context, d *coreDependencies, closers *closeFuncs) *pg.DB {
+	pg.UseLogger(d.logger.New("PG"))
+
 	// TODO: restore from snapshots
 
 	db, err := d.dbOpener(ctx, d.cfg.DB.DBName, d.cfg.DB.MaxConns)

node/consensus/engine_test.go

@@ -175,6 +175,11 @@ func verifyStatus(_ *testing.T, val *ConsensusEngine, status Status, height int6
 	return nil
 }
 
+func TestMain(m *testing.M) {
+	pg.UseLogger(log.New(log.WithName("DBS"), log.WithFormat(log.FormatUnstructured)))
+	m.Run()
+}
+
 func TestValidatorStateMachine(t *testing.T) {
 	// t.Parallel()
 	type action struct {

node/node_live_test.go

@@ -41,6 +41,11 @@ var defaultGenesisParams = &consensus.GenesisParams{
 	},
 }
 
+func TestMain(m *testing.M) {
+	pg.UseLogger(log.New(log.WithName("DBS"), log.WithFormat(log.FormatUnstructured)))
+	m.Run()
+}
+
 func TestSingleNodeMocknet(t *testing.T) {
 	if testing.Short() {
 		t.Skip()
@@ -54,8 +59,6 @@ func TestSingleNodeMocknet(t *testing.T) {
 	bs1 := memstore.NewMemBS()
 	mp1 := mempool.New()
 
-	pg.UseLogger(log.New(log.WithName("DBS")))
-
 	db1 := initDB(t, "5432", "kwil_test_db")
 
 	root1 := t.TempDir()
@@ -155,8 +158,6 @@ func TestDualNodeMocknet(t *testing.T) {
 	bs1 := memstore.NewMemBS()
 	mp1 := mempool.New()
 
-	pg.UseLogger(log.New(log.WithName("DBS")))
-
 	db1 := initDB(t, "5432", "kwil_test_db")
 	func() {
 		ctx := context.Background()
@@ -332,10 +333,9 @@ func initDB(t *testing.T, port, dbName string) *pg.DB {
 }
 
 func cleanupDB(db *pg.DB) {
-	ctx := context.Background()
-	db.AutoCommit(true)
-	defer db.AutoCommit(false)
 	defer db.Close()
+	db.AutoCommit(true)
+	ctx := context.Background()
 	db.Execute(ctx, `DROP SCHEMA IF EXISTS kwild_chain CASCADE;`)
 	db.Execute(ctx, `DROP SCHEMA IF EXISTS kwild_internal CASCADE;`)
 }

node/pg/db.go

@@ -266,17 +266,34 @@ func (db *DB) EnsureFullReplicaIdentityDatasets(ctx context.Context) error {
 func (db *DB) Close() error {
 	db.cancel(nil)
 	db.repl.stop()
-	if db.tx != nil {
-		// This is a bug, so we are going to panic so the issue is not easily
-		// ignored, but we will rollback the tx so we don't hang or leak
-		// postgresql resources.
-		db.tx.Rollback(context.Background())
-		db.tx = nil
+	if db.tx == nil {
+		return db.pool.Close()
+	}
+
+	// This is a bug, so we are going to panic so the issue is not easily
+	// ignored, but we will rollback the tx so we don't hang or leak
+	// postgresql resources.
+
+	if db.txid != "" {
+		logger.Warnln("Rolling back PREPARED transaction", db.txid)
+		// With PREPARE TRANSACTION already done, rollback the prepared transaction.
+		sqlRollback := fmt.Sprintf(`ROLLBACK PREPARED '%s'`, db.txid)
+		if _, err := db.tx.Exec(context.Background(), sqlRollback); err != nil {
+			return fmt.Errorf("ROLLBACK PREPARED failed: %v", err)
+		}
+		// We don't use Rollback, which normally releases automatically.
+		if rel, ok := db.tx.(releaser); ok {
+			rel.Release()
+		}
 		db.txid = ""
-		db.pool.Close()
-		panic("Closed the DB with an active transaction!")
+	} else { // otherwise regular rollback
+		logger.Warnln("Rolling back regular transaction")
+		db.tx.Rollback(context.Background())
 	}
-	return db.pool.Close()
+
+	db.tx = nil
+	db.pool.Close()
+	panic("Closed the DB with an active transaction!")
 }
 
 // Done allows higher level systems to monitor the state of the DB backend