fix memory leaks in PrepareStatementDB

prepare_stmt.go

@@ -17,18 +17,16 @@ type Stmt struct {
 }
 
 type PreparedStmtDB struct {
-	Stmts       map[string]*Stmt
-	PreparedSQL []string
-	Mux         *sync.RWMutex
+	Stmts map[string]*Stmt
+	Mux   *sync.RWMutex
 	ConnPool
 }
 
 func NewPreparedStmtDB(connPool ConnPool) *PreparedStmtDB {
 	return &PreparedStmtDB{
-		ConnPool:    connPool,
-		Stmts:       make(map[string]*Stmt),
-		Mux:         &sync.RWMutex{},
-		PreparedSQL: make([]string, 0, 100),
+		ConnPool: connPool,
+		Stmts:    make(map[string]*Stmt),
+		Mux:      &sync.RWMutex{},
 	}
 }
 
@@ -48,12 +46,18 @@ func (db *PreparedStmtDB) Close() {
 	db.Mux.Lock()
 	defer db.Mux.Unlock()
 
-	for _, query := range db.PreparedSQL {
-		if stmt, ok := db.Stmts[query]; ok {
-			delete(db.Stmts, query)
-			go stmt.Close()
-		}
+	for _, stmt := range db.Stmts {
+		go func(s *Stmt) {
+			// make sure the stmt must finish preparation first
+			<-s.prepared
+			if s.Stmt != nil {
+				_ = s.Close()
+			}
+		}(stmt)
 	}
+	// I think setting it to nil should be better, as there should be no one
+	// writing into it after calling Close, but leave it for compatibility and safety
+	db.Stmts = make(map[string]*Stmt)
 }
 
 func (sdb *PreparedStmtDB) Reset() {
@@ -63,7 +67,6 @@ func (sdb *PreparedStmtDB) Reset() {
 	for _, stmt := range sdb.Stmts {
 		go stmt.Close()
 	}
-	sdb.PreparedSQL = make([]string, 0, 100)
 	sdb.Stmts = make(map[string]*Stmt)
 }
 
@@ -118,7 +121,6 @@ func (db *PreparedStmtDB) prepare(ctx context.Context, conn ConnPool, isTransact
 
 	db.Mux.Lock()
 	cacheStmt.Stmt = stmt
-	db.PreparedSQL = append(db.PreparedSQL, query)
 	db.Mux.Unlock()
 
 	return cacheStmt, nil