Skip to content

Commit 9e7a0a8

Browse files
ncrucesncruces
committed
Improved arg reuse.
1 parent 26adda4 commit 9e7a0a8

File tree

5 files changed

+51
-47
lines changed

5 files changed

+51
-47
lines changed

‎const.go

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -11,10 +11,9 @@ const (
1111
_ROW = 100 /* sqlite3_step() has another row ready */
1212
_DONE = 101 /* sqlite3_step() has finished executing */
1313

14-
_MAX_NAME = 1e6 // Self-imposed limit for most NUL terminated strings.
15-
_MAX_LENGTH = 1e9
16-
_MAX_SQL_LENGTH = 1e9
17-
_MAX_FUNCTION_ARG = 100
14+
_MAX_NAME = 1e6 // Self-imposed limit for most NUL terminated strings.
15+
_MAX_LENGTH = 1e9
16+
_MAX_SQL_LENGTH = 1e9
1817

1918
ptrlen = util.PtrLen
2019
intlen = util.IntLen

‎ext/bloom/bloom.go

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -268,13 +268,13 @@ func (b *bloom) Open() (sqlite3.VTabCursor, error) {
268268

269269
type cursor struct {
270270
*bloom
271-
arg *sqlite3.Value
271+
arg sqlite3.Value
272272
eof bool
273273
}
274274

275275
func (c *cursor) Filter(idxNum int, idxStr string, arg ...sqlite3.Value) error {
276276
c.eof = false
277-
c.arg = &arg[0]
277+
c.arg = arg[0]
278278
blob := arg[0].RawBlob()
279279

280280
f, err := c.db.OpenBlob(c.schema, c.storage, "data", 1, false)
@@ -312,7 +312,7 @@ func (c *cursor) Column(ctx sqlite3.Context, n int) error {
312312
case 0:
313313
ctx.ResultBool(true)
314314
case 1:
315-
ctx.ResultValue(*c.arg)
315+
ctx.ResultValue(c.arg)
316316
}
317317
return nil
318318
}

‎ext/statement/stmt.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -159,8 +159,6 @@ func (c *cursor) Close() error {
159159
}
160160

161161
func (c *cursor) Filter(idxNum int, idxStr string, arg ...sqlite3.Value) error {
162-
c.arg = arg
163-
c.rowID = 0
164162
err := errors.Join(
165163
c.stmt.Reset(),
166164
c.stmt.ClearBindings())
@@ -187,6 +185,8 @@ func (c *cursor) Filter(idxNum int, idxStr string, arg ...sqlite3.Value) error {
187185
return err
188186
}
189187
}
188+
c.arg = append(c.arg[:0], arg...)
189+
c.rowID = 0
190190
return c.Next()
191191
}
192192

‎func.go

Lines changed: 30 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ import (
55
"io"
66
"iter"
77
"sync"
8+
"sync/atomic"
89

910
"github.com/tetratelabs/wazero/api"
1011

@@ -196,21 +197,19 @@ func compareCallback(ctx context.Context, mod api.Module, pApp ptr_t, nKey1 int3
196197
}
197198

198199
func funcCallback(ctx context.Context, mod api.Module, pCtx, pApp ptr_t, nArg int32, pArg ptr_t) {
199-
args := getFuncArgs()
200-
defer putFuncArgs(args)
201200
db := ctx.Value(connKey{}).(*Conn)
201+
args := callbackArgs(db, nArg, pArg)
202+
defer returnArgs(args)
202203
fn := util.GetHandle(db.ctx, pApp).(ScalarFunction)
203-
callbackArgs(db, args[:nArg], pArg)
204-
fn(Context{db, pCtx}, args[:nArg]...)
204+
fn(Context{db, pCtx}, *args...)
205205
}
206206

207207
func stepCallback(ctx context.Context, mod api.Module, pCtx, pAgg, pApp ptr_t, nArg int32, pArg ptr_t) {
208-
args := getFuncArgs()
209-
defer putFuncArgs(args)
210208
db := ctx.Value(connKey{}).(*Conn)
211-
callbackArgs(db, args[:nArg], pArg)
209+
args := callbackArgs(db, nArg, pArg)
210+
defer returnArgs(args)
212211
fn, _ := callbackAggregate(db, pAgg, pApp)
213-
fn.Step(Context{db, pCtx}, args[:nArg]...)
212+
fn.Step(Context{db, pCtx}, *args...)
214213
}
215214

216215
func valueCallback(ctx context.Context, mod api.Module, pCtx, pAgg, pApp ptr_t, final int32) {
@@ -234,12 +233,11 @@ func valueCallback(ctx context.Context, mod api.Module, pCtx, pAgg, pApp ptr_t,
234233
}
235234

236235
func inverseCallback(ctx context.Context, mod api.Module, pCtx, pAgg ptr_t, nArg int32, pArg ptr_t) {
237-
args := getFuncArgs()
238-
defer putFuncArgs(args)
239236
db := ctx.Value(connKey{}).(*Conn)
240-
callbackArgs(db, args[:nArg], pArg)
237+
args := callbackArgs(db, nArg, pArg)
238+
defer returnArgs(args)
241239
fn := util.GetHandle(db.ctx, pAgg).(WindowFunction)
242-
fn.Inverse(Context{db, pCtx}, args[:nArg]...)
240+
fn.Inverse(Context{db, pCtx}, *args...)
243241
}
244242

245243
func callbackAggregate(db *Conn, pAgg, pApp ptr_t) (AggregateFunction, ptr_t) {
@@ -258,28 +256,31 @@ func callbackAggregate(db *Conn, pAgg, pApp ptr_t) (AggregateFunction, ptr_t) {
258256
return fn, 0
259257
}
260258

261-
func callbackArgs(db *Conn, arg []Value, pArg ptr_t) {
262-
for i := range arg {
263-
arg[i] = Value{
259+
var (
260+
valueArgsPool sync.Pool
261+
valueArgsLen atomic.Int32
262+
)
263+
264+
func callbackArgs(db *Conn, nArg int32, pArg ptr_t) *[]Value {
265+
arg, ok := valueArgsPool.Get().(*[]Value)
266+
if !ok || cap(*arg) < int(nArg) {
267+
max := valueArgsLen.Or(nArg) | nArg
268+
lst := make([]Value, max)
269+
arg = &lst
270+
}
271+
lst := (*arg)[:nArg]
272+
for i := range lst {
273+
lst[i] = Value{
264274
c: db,
265275
handle: util.Read32[ptr_t](db.mod, pArg+ptr_t(i)*ptrlen),
266276
}
267277
}
278+
*arg = lst
279+
return arg
268280
}
269281

270-
var funcArgsPool sync.Pool
271-
272-
func putFuncArgs(p *[_MAX_FUNCTION_ARG]Value) {
273-
clear(p[:])
274-
funcArgsPool.Put(p)
275-
}
276-
277-
func getFuncArgs() *[_MAX_FUNCTION_ARG]Value {
278-
if p := funcArgsPool.Get(); p == nil {
279-
return new([_MAX_FUNCTION_ARG]Value)
280-
} else {
281-
return p.(*[_MAX_FUNCTION_ARG]Value)
282-
}
282+
func returnArgs(p *[]Value) {
283+
valueArgsPool.Put(p)
283284
}
284285

285286
type aggregateFunc struct {
@@ -291,7 +292,7 @@ type aggregateFunc struct {
291292

292293
func (a *aggregateFunc) Step(ctx Context, arg ...Value) {
293294
a.ctx = ctx
294-
a.arg = arg
295+
a.arg = append(a.arg[:0], arg...)
295296
if _, more := a.next(); !more {
296297
a.stop()
297298
}

‎vtab.go

Lines changed: 13 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -162,6 +162,7 @@ type VTabDestroyer interface {
162162
}
163163

164164
// A VTabUpdater allows a virtual table to be updated.
165+
// Implementations must not retain arg.
165166
type VTabUpdater interface {
166167
VTab
167168
// https://sqlite.org/vtab.html#xupdate
@@ -241,6 +242,7 @@ type VTabSavepointer interface {
241242
// to loop through the virtual table.
242243
// A VTabCursor may optionally implement
243244
// [io.Closer] to free resources.
245+
// Implementations of Filter must not retain arg.
244246
//
245247
// https://sqlite.org/c3ref/vtab_cursor.html
246248
type VTabCursor interface {
@@ -489,12 +491,12 @@ func vtabBestIndexCallback(ctx context.Context, mod api.Module, pVTab, pIdxInfo
489491
}
490492

491493
func vtabUpdateCallback(ctx context.Context, mod api.Module, pVTab ptr_t, nArg int32, pArg, pRowID ptr_t) res_t {
492-
vtab := vtabGetHandle(ctx, mod, pVTab).(VTabUpdater)
493-
494494
db := ctx.Value(connKey{}).(*Conn)
495-
args := make([]Value, nArg)
496-
callbackArgs(db, args, pArg)
497-
rowID, err := vtab.Update(args...)
495+
args := callbackArgs(db, nArg, pArg)
496+
defer returnArgs(args)
497+
498+
vtab := vtabGetHandle(ctx, mod, pVTab).(VTabUpdater)
499+
rowID, err := vtab.Update(*args...)
498500
if err == nil {
499501
util.Write64(mod, pRowID, rowID)
500502
}
@@ -593,15 +595,17 @@ func cursorCloseCallback(ctx context.Context, mod api.Module, pCur ptr_t) res_t
593595
}
594596

595597
func cursorFilterCallback(ctx context.Context, mod api.Module, pCur ptr_t, idxNum int32, idxStr ptr_t, nArg int32, pArg ptr_t) res_t {
596-
cursor := vtabGetHandle(ctx, mod, pCur).(VTabCursor)
597598
db := ctx.Value(connKey{}).(*Conn)
598-
args := make([]Value, nArg)
599-
callbackArgs(db, args, pArg)
599+
args := callbackArgs(db, nArg, pArg)
600+
defer returnArgs(args)
601+
600602
var idxName string
601603
if idxStr != 0 {
602604
idxName = util.ReadString(mod, idxStr, _MAX_LENGTH)
603605
}
604-
err := cursor.Filter(int(idxNum), idxName, args...)
606+
607+
cursor := vtabGetHandle(ctx, mod, pCur).(VTabCursor)
608+
err := cursor.Filter(int(idxNum), idxName, *args...)
605609
return vtabError(ctx, mod, pCur, _CURSOR_ERROR, err)
606610
}
607611

0 commit comments

Comments
 (0)