Skip to content

Commit

Permalink
Fix exponential memory allocation in SQLiteConn.Exec
Browse files Browse the repository at this point in the history
This commit fixes SQLiteConn.Exec to not duplicate the provided SQL
statement and to pass the Go string to libsqlite3 instead of creating
a C string copy and passing that. The issue with the previous
implementation was that it would create a new C string copy of the
original Go string each time it stepped through a SQL statement. This
led to exponential memory usage when Exec'ing an argument that contained
multiple SQL statements (which is common when initializing a database
from a SQL dump).

This commit is a slimmed down version of PR mattn#1133:
  mattn#1133

```
goos: darwin
goarch: arm64
pkg: github.com/mattn/go-sqlite3
cpu: Apple M1 Max
                           │    b.txt     │                n.txt                │
                           │    sec/op    │   sec/op     vs base                │
Suite/BenchmarkExec-10        1.278µ ± 1%   1.022µ ± 2%  -19.99% (p=0.000 n=10)
Suite/BenchmarkExecStep-10   1762.9µ ± 0%   900.9µ ± 1%  -48.90% (p=0.000 n=10)
geomean                       47.47µ        30.35µ       -36.06%

                           │     b.txt      │                n.txt                 │
                           │      B/op      │     B/op      vs base                │
Suite/BenchmarkExec-10           128.0 ± 0%     120.0 ± 0%   -6.25% (p=0.000 n=10)
Suite/BenchmarkExecStep-10   5279.94Ki ± 0%   31.34Ki ± 0%  -99.41% (p=0.000 n=10)
geomean                        25.69Ki        1.916Ki       -92.54%

                           │    b.txt    │                n.txt                │
                           │  allocs/op  │  allocs/op   vs base                │
Suite/BenchmarkExec-10        7.000 ± 0%    6.000 ± 0%  -14.29% (p=0.000 n=10)
Suite/BenchmarkExecStep-10   7.000k ± 0%   3.003k ± 0%  -57.10% (p=0.000 n=10)
geomean                       221.4         134.2       -39.36%
```
  • Loading branch information
charlievieth committed Nov 7, 2024
1 parent b3e6ac1 commit 6915212
Show file tree
Hide file tree
Showing 4 changed files with 106 additions and 22 deletions.
77 changes: 55 additions & 22 deletions sqlite3.go
Original file line number Diff line number Diff line change
Expand Up @@ -137,6 +137,22 @@ _sqlite3_prepare_v2_internal(sqlite3 *db, const char *zSql, int nBytes, sqlite3_
}
#endif
static int _sqlite3_prepare_v2(sqlite3 *db, const char *zSql, int nBytes, sqlite3_stmt **ppStmt, int *oBytes) {
const char *tail = NULL;
int rv = _sqlite3_prepare_v2_internal(db, zSql, nBytes, ppStmt, &tail);
if (rv != SQLITE_OK) {
return rv;
}
if (tail == NULL) {
return rv; // NB: this should not happen
}
// Set oBytes to the number of bytes consumed instead of using the **pzTail
// out param since that requires storing a Go pointer in a C pointer, which
// is not allowed by CGO and will cause runtime.cgoCheckPointer to fail.
*oBytes = tail - zSql;
return rv;
}
void _sqlite3_result_text(sqlite3_context* ctx, const char* s) {
sqlite3_result_text(ctx, s, -1, &free);
}
Expand Down Expand Up @@ -858,51 +874,61 @@ func (c *SQLiteConn) Exec(query string, args []driver.Value) (driver.Result, err
}

func (c *SQLiteConn) exec(ctx context.Context, query string, args []driver.NamedValue) (driver.Result, error) {
start := 0
var (
stmtArgs []driver.NamedValue
start int
s SQLiteStmt // escapes to the heap so reuse it
sz C.int // number of query bytes consumed: escapes to the heap
)
query = strings.TrimSpace(query)
for {
s, err := c.prepare(ctx, query)
if err != nil {
return nil, err
s = SQLiteStmt{c: c} // reset
sz = 0
rv := C._sqlite3_prepare_v2(c.db, (*C.char)(unsafe.Pointer(stringData(query))),
C.int(len(query)), &s.s, &sz)
if rv != C.SQLITE_OK {
return nil, c.lastError()
}
query = strings.TrimSpace(query[sz:])

var res driver.Result
if s.(*SQLiteStmt).s != nil {
stmtArgs := make([]driver.NamedValue, 0, len(args))
if s.s != nil {
na := s.NumInput()
if len(args)-start < na {
s.Close()
s.finalize()
return nil, fmt.Errorf("not enough args to execute query: want %d got %d", na, len(args))
}
// consume the number of arguments used in the current
// statement and append all named arguments not
// contained therein
if len(args[start:start+na]) > 0 {
stmtArgs = append(stmtArgs, args[start:start+na]...)
for i := range args {
if (i < start || i >= na) && args[i].Name != "" {
stmtArgs = append(stmtArgs, args[i])
}
}
for i := range stmtArgs {
stmtArgs[i].Ordinal = i + 1
if stmtArgs == nil {
stmtArgs = make([]driver.NamedValue, 0, na)
}
stmtArgs = append(stmtArgs[:0], args[start:start+na]...)
for i := range args {
if (i < start || i >= na) && args[i].Name != "" {
stmtArgs = append(stmtArgs, args[i])
}
}
res, err = s.(*SQLiteStmt).exec(ctx, stmtArgs)
for i := range stmtArgs {
stmtArgs[i].Ordinal = i + 1
}
var err error
res, err = s.exec(ctx, stmtArgs)
if err != nil && err != driver.ErrSkip {
s.Close()
s.finalize()
return nil, err
}
start += na
}
tail := s.(*SQLiteStmt).t
s.Close()
if tail == "" {
s.finalize()
if len(query) == 0 {
if res == nil {
// https://github.com/mattn/go-sqlite3/issues/963
res = &SQLiteResult{0, 0}
}
return res, nil
}
query = tail
}
}

Expand Down Expand Up @@ -1914,6 +1940,13 @@ func (s *SQLiteStmt) Close() error {
return nil
}

func (s *SQLiteStmt) finalize() {
if s.s != nil {
C.sqlite3_finalize(s.s)
s.s = nil
}
}

// NumInput return a number of parameters.
func (s *SQLiteStmt) NumInput() int {
return int(C.sqlite3_bind_parameter_count(s.s))
Expand Down
11 changes: 11 additions & 0 deletions sqlite3_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -2106,6 +2106,7 @@ var tests = []testing.InternalTest{

var benchmarks = []testing.InternalBenchmark{
{Name: "BenchmarkExec", F: benchmarkExec},
{Name: "BenchmarkExecStep", F: benchmarkExecStep},
{Name: "BenchmarkQuery", F: benchmarkQuery},
{Name: "BenchmarkParams", F: benchmarkParams},
{Name: "BenchmarkStmt", F: benchmarkStmt},
Expand Down Expand Up @@ -2465,6 +2466,16 @@ func benchmarkExec(b *testing.B) {
}
}

var largeSelectStmt = strings.Repeat("select 1;\n", 1_000)

func benchmarkExecStep(b *testing.B) {
for n := 0; n < b.N; n++ {
if _, err := db.Exec(largeSelectStmt); err != nil {
b.Fatal(err)
}
}
}

// benchmarkQuery is benchmark for query
func benchmarkQuery(b *testing.B) {
for i := 0; i < b.N; i++ {
Expand Down
17 changes: 17 additions & 0 deletions unsafe_go120.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,17 @@
//go:build !go1.21
// +build !go1.21

package sqlite3

import "unsafe"

// stringData is a safe version of unsafe.StringData that handles empty strings.
func stringData(s string) *byte {
if len(s) != 0 {
b := *(*[]byte)(unsafe.Pointer(&s))
return &b[0]
}
// The return value of unsafe.StringData
// is unspecified if the string is empty.
return &placeHolder[0]
}
23 changes: 23 additions & 0 deletions unsafe_go121.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,23 @@
//go:build go1.21
// +build go1.21

// The unsafe.StringData function was made available in Go 1.20 but it
// was not until Go 1.21 that Go was changed to interpret the Go version
// in go.mod (1.19 as of writing this) as the minimum version required
// instead of the exact version.
//
// See: https://github.com/golang/go/issues/59033

package sqlite3

import "unsafe"

// stringData is a safe version of unsafe.StringData that handles empty strings.
func stringData(s string) *byte {
if len(s) != 0 {
return unsafe.StringData(s)
}
// The return value of unsafe.StringData
// is unspecified if the string is empty.
return &placeHolder[0]
}

0 comments on commit 6915212

Please sign in to comment.