Reduce SQL sanitizer allocations #2136

ninedraft · 2024-10-01T14:43:53Z

#2124

Result:

Main optimizations:

extensive usage of sync.Pool for byte buffers, lexers and parsed query structs
append-style string formatters for int64, float64 and time.Time + bytes.Buffer.AvailableBuffer
rework of QuoteString and QuoteBytes to append-style (with tests for backwards compatibility)

Misc changes:

benchmarks for Query.Sanitize and SanitizeSQL functions
a tiny script for generation of benchmark reports for selected commits and diff (using benchstat)
fuzzing of QuoteString and QuoteBytes (I did'n find any problems for 1h of fuzzing, but you can never be sure for 100%)

Since optimization is an extremely hard problem, I think it's worth checking some more benchmarks.

I would be very grateful for your opinion on this and recommendations/advice, @jackc @vtolstov

ninedraft · 2024-10-01T15:01:23Z

benchmark diffs for concrete optimisations

goos: darwin
goarch: arm64
pkg: github.com/jackc/pgx/v5/internal/sanitize
cpu: Apple M1
              │ benchmarks/0_base_case.bench │     benchmarks/1_buf_pool.bench     │ benchmarks/2_append_AvailableBuffer.bench │    benchmarks/3_quoteBytes.bench    │   benchmarks/4_quoteString.bench    │ benchmarks/5_add_lexer_and_query_pools.bench │ benchmarks/6_drop_too_large_values_from_memory_pools.bench │
              │            sec/op            │   sec/op     vs base                │      sec/op        vs base                │   sec/op     vs base                │   sec/op     vs base                │        sec/op         vs base                │               sec/op                vs base                │
Sanitize-8                       718.2n ± 1%   578.8n ± 1%  -19.41% (p=0.000 n=10)         439.9n ± 0%  -38.74% (p=0.000 n=10)   413.6n ± 4%  -42.42% (p=0.000 n=10)   397.1n ± 1%  -44.72% (p=0.000 n=10)            403.6n ± 1%  -43.81% (p=0.000 n=10)                          400.8n ± 2%  -44.20% (p=0.000 n=10)
SanitizeSQL-8                    2.089µ ± 0%   1.956µ ± 0%   -6.37% (p=0.000 n=10)         1.828µ ± 0%  -12.49% (p=0.000 n=10)   1.812µ ± 1%  -13.28% (p=0.000 n=10)   1.789µ ± 1%  -14.36% (p=0.000 n=10)            1.670µ ± 0%  -20.06% (p=0.000 n=10)                          1.673µ ± 0%  -19.91% (p=0.000 n=10)
geomean                          1.225µ        1.064µ       -13.13%                        896.8n       -26.79%                  865.5n       -29.34%                  842.8n       -31.19%                           820.9n       -32.98%                                         818.8n       -33.15%

              │ benchmarks/0_base_case.bench │     benchmarks/1_buf_pool.bench     │ benchmarks/2_append_AvailableBuffer.bench │    benchmarks/3_quoteBytes.bench    │   benchmarks/4_quoteString.bench    │ benchmarks/5_add_lexer_and_query_pools.bench │ benchmarks/6_drop_too_large_values_from_memory_pools.bench │
              │             B/op             │    B/op      vs base                │       B/op         vs base                │    B/op      vs base                │    B/op      vs base                │         B/op          vs base                │                B/op                 vs base                │
Sanitize-8                       1488.0 ± 0%    528.0 ± 0%  -64.52% (p=0.000 n=10)          472.0 ± 0%  -68.28% (p=0.000 n=10)    456.0 ± 0%  -69.35% (p=0.000 n=10)    424.0 ± 0%  -71.51% (p=0.000 n=10)             424.0 ± 0%  -71.51% (p=0.000 n=10)                           424.0 ± 0%  -71.51% (p=0.000 n=10)
SanitizeSQL-8                    2216.0 ± 0%   1256.0 ± 0%  -43.32% (p=0.000 n=10)         1200.0 ± 0%  -45.85% (p=0.000 n=10)   1184.0 ± 0%  -46.57% (p=0.000 n=10)   1152.0 ± 0%  -48.01% (p=0.000 n=10)             552.0 ± 0%  -75.09% (p=0.000 n=10)                           552.0 ± 0%  -75.09% (p=0.000 n=10)
geomean                         1.773Ki         814.4       -55.15%                         752.6       -58.55%                   734.8       -59.54%                   698.9       -61.51%                            483.8       -73.36%                                          483.8       -73.36%

              │ benchmarks/0_base_case.bench │    benchmarks/1_buf_pool.bench     │ benchmarks/2_append_AvailableBuffer.bench │   benchmarks/3_quoteBytes.bench    │   benchmarks/4_quoteString.bench   │ benchmarks/5_add_lexer_and_query_pools.bench │ benchmarks/6_drop_too_large_values_from_memory_pools.bench │
              │          allocs/op           │ allocs/op   vs base                │     allocs/op      vs base                │ allocs/op   vs base                │ allocs/op   vs base                │      allocs/op        vs base                │             allocs/op               vs base                │
Sanitize-8                       11.000 ± 0%   7.000 ± 0%  -36.36% (p=0.000 n=10)          4.000 ± 0%  -63.64% (p=0.000 n=10)   3.000 ± 0%  -72.73% (p=0.000 n=10)   2.000 ± 0%  -81.82% (p=0.000 n=10)             2.000 ± 0%  -81.82% (p=0.000 n=10)                           2.000 ± 0%  -81.82% (p=0.000 n=10)
SanitizeSQL-8                     26.00 ± 0%   22.00 ± 0%  -15.38% (p=0.000 n=10)          19.00 ± 0%  -26.92% (p=0.000 n=10)   18.00 ± 0%  -30.77% (p=0.000 n=10)   17.00 ± 0%  -34.62% (p=0.000 n=10)             10.00 ± 0%  -61.54% (p=0.000 n=10)                           10.00 ± 0%  -61.54% (p=0.000 n=10)
geomean                           16.91        12.41       -26.62%                         8.718       -48.45%                  7.348       -56.55%                  5.831       -65.52%                            4.472       -73.56%                                          4.472       -73.56%

jackc · 2024-10-05T14:50:38Z

LGTM. But this is obviously a very security critical part of the code, so I'd like if we can get some more eyes on this before merging.

vtolstov · 2024-10-06T13:05:58Z

lgtm, i'm try to check on my hot path in next few days.

vtolstov · 2024-10-10T21:05:00Z

In my tests i don't saw any issues.

ninedraft · 2024-10-15T17:03:52Z

@jackc

But this is obviously a very security critical part of the code, so I'd like if we can get some more eyes on this before merging.

It would be very much appreciated if you could suggest someone I can tag on this issue. I'm also in the process of writing more tests for SQL injection + more fuzzing

jackc · 2024-10-18T21:45:56Z

It would be very much appreciated if you could suggest someone I can tag on this issue.

I wish I could. Unfortunately, I don't know of anyone.

I'm also in the process of writing more tests for SQL injection + more fuzzing

👍 It's been a couple weeks since I reviewed the code, so I can review it again with fresh eyes now. It's not quite as good as multiple reviewers, but at least it will be multiple reviews.

I'll wait until you add the additional tests.

sean-

See the inline suggested changes. These small set of optimizations reduced the sec/op from 607.9n down to 604.9n (-0.49%). I can push this as a different PR or these changes can be incorporated into this branch.

$ bash benchmmark.sh 2ec900454bfe65daa9648488e93f7627c26b810c 82642726914a8b054ca123fd87c4d984da6d78eb 431e11b61c809c2373128ecf63ed48cf8bdf4dd4 71c3b107187b02ea44dbc7d38e931115ca7286c7
$ benchstat benchmarks/*.bench
goos: darwin
goarch: arm64
pkg: github.com/jackc/pgx/v5/internal/sanitize
cpu: Apple M3 Pro
               │ benchmarks/1_fix_preallocations_of_quoted_string.bench │ benchmarks/2_optimize_QuoteBytes.bench │ benchmarks/2_rework_QuoteString_to_iterate_over_bytes.bench │ benchmarks/3_optimize_QuoteBytes.bench │
               │                         sec/op                         │     sec/op       vs base               │                sec/op                 vs base               │     sec/op       vs base               │
Sanitize-12                                                 307.1n ± 1%       300.6n ± 2%  -2.10% (p=0.001 n=10)                            305.3n ± 1%  -0.57% (p=0.015 n=10)       304.2n ± 1%  -0.93% (p=0.003 n=10)
SanitizeSQL-12                                              1.204µ ± 2%       1.207µ ± 1%       ~ (p=0.100 n=10)                            1.204µ ± 2%       ~ (p=0.697 n=10)       1.203µ ± 2%       ~ (p=0.898 n=10)
geomean                                                     607.9n            602.3n       -0.91%                                           606.3n       -0.26%                      604.9n       -0.49%

               │ benchmarks/1_fix_preallocations_of_quoted_string.bench │ benchmarks/2_optimize_QuoteBytes.bench │ benchmarks/2_rework_QuoteString_to_iterate_over_bytes.bench │ benchmarks/3_optimize_QuoteBytes.bench │
               │                          B/op                          │     B/op       vs base                 │                B/op                 vs base                 │     B/op       vs base                 │
Sanitize-12                                                  424.0 ± 0%      424.0 ± 0%       ~ (p=1.000 n=10) ¹                           424.0 ± 0%       ~ (p=1.000 n=10) ¹      424.0 ± 0%       ~ (p=1.000 n=10) ¹
SanitizeSQL-12                                               552.0 ± 0%      552.0 ± 0%       ~ (p=1.000 n=10) ¹                           552.0 ± 0%       ~ (p=1.000 n=10) ¹      552.0 ± 0%       ~ (p=1.000 n=10) ¹
geomean                                                      483.8           483.8       +0.00%                                            483.8       +0.00%                       483.8       +0.00%
¹ all samples are equal

               │ benchmarks/1_fix_preallocations_of_quoted_string.bench │ benchmarks/2_optimize_QuoteBytes.bench │ benchmarks/2_rework_QuoteString_to_iterate_over_bytes.bench │ benchmarks/3_optimize_QuoteBytes.bench │
               │                       allocs/op                        │   allocs/op    vs base                 │             allocs/op               vs base                 │   allocs/op    vs base                 │
Sanitize-12                                                  2.000 ± 0%      2.000 ± 0%       ~ (p=1.000 n=10) ¹                           2.000 ± 0%       ~ (p=1.000 n=10) ¹      2.000 ± 0%       ~ (p=1.000 n=10) ¹
SanitizeSQL-12                                               10.00 ± 0%      10.00 ± 0%       ~ (p=1.000 n=10) ¹                           10.00 ± 0%       ~ (p=1.000 n=10) ¹      10.00 ± 0%       ~ (p=1.000 n=10) ¹
geomean                                                      4.472           4.472       +0.00%                                            4.472       +0.00%                       4.472       +0.00%
¹ all samples are equal

sean- · 2024-10-21T15:22:53Z

internal/sanitize/benchmmark.sh

+    }
+
+    # Sanitized commmit message
+    commit_message=$(git log -1 --pretty=format:"%s" | tr ' ' '_')


This needs to escape /:

commit_message=$(git log -1 --pretty=format:"%s" | tr -c '[:alnum:]-_' '_')

sean- · 2024-10-21T15:23:51Z

internal/sanitize/benchmmark.sh

+    bench_files+=("$bench_file")
+done
+
+benchstat "${bench_files[@]}"


Can you prefix with a small comment: # go install golang.org/x/perf/cmd/benchstat@latest

sean- · 2024-10-21T15:26:50Z

internal/sanitize/sanitize.go

+
+	dst = append(dst, quote...)
+
+	return dst
 }


This is purely a style nit, but I don't like reslicing for these types of functions because it's not idiomatic and hard to follow. I took the above QuoteString() and replaced it with something that uses an iterator:

func QuoteString(dst []byte, str string) []byte { const quote = '\'' // Preallocate space for the worst case scenario dst = slices.Grow(dst, len(str)*2+2) // Add opening quote dst = append(dst, quote) // Iterate through the string without allocating for i := 0; i < len(str); i++ { if str[i] == quote { dst = append(dst, quote, quote) } else { dst = append(dst, str[i]) } } // Add closing quote dst = append(dst, quote) return dst }

sean- · 2024-10-21T15:27:39Z

internal/sanitize/sanitize.go

+	dst = append(dst, p...)
+
+	dst = append(dst, `'`...)
+	return dst
 }


I was able to measure an improvement by optimizing this function:

func QuoteBytes(dst, buf []byte) []byte { if len(buf) == 0 { return append(dst, `'\x'`...) } // Calculate required length requiredLen := 3 + hex.EncodedLen(len(buf)) + 1 // Ensure dst has enough capacity if cap(dst)-len(dst) < requiredLen { newDst := make([]byte, len(dst), len(dst)+requiredLen) copy(newDst, dst) dst = newDst } // Record original length and extend slice origLen := len(dst) dst = dst[:origLen+requiredLen] // Add prefix dst[origLen] = '\'' dst[origLen+1] = '\\' dst[origLen+2] = 'x' // Encode bytes directly into dst hex.Encode(dst[origLen+3:len(dst)-1], buf) // Add suffix dst[len(dst)-1] = '\'' return dst }

sean-

See the inline suggested changes. These small set of optimizations reduced the sec/op from 607.9n down to 604.9n (-0.49%). I can push this as a different PR or these changes can be incorporated into this branch.

$ bash benchmmark.sh 2ec900454bfe65daa9648488e93f7627c26b810c 82642726914a8b054ca123fd87c4d984da6d78eb 431e11b61c809c2373128ecf63ed48cf8bdf4dd4 71c3b107187b02ea44dbc7d38e931115ca7286c7
$ benchstat benchmarks/*.bench
goos: darwin
goarch: arm64
pkg: github.com/jackc/pgx/v5/internal/sanitize
cpu: Apple M3 Pro
               │ benchmarks/1_fix_preallocations_of_quoted_string.bench │ benchmarks/2_optimize_QuoteBytes.bench │ benchmarks/2_rework_QuoteString_to_iterate_over_bytes.bench │ benchmarks/3_optimize_QuoteBytes.bench │
               │                         sec/op                         │     sec/op       vs base               │                sec/op                 vs base               │     sec/op       vs base               │
Sanitize-12                                                 307.1n ± 1%       300.6n ± 2%  -2.10% (p=0.001 n=10)                            305.3n ± 1%  -0.57% (p=0.015 n=10)       304.2n ± 1%  -0.93% (p=0.003 n=10)
SanitizeSQL-12                                              1.204µ ± 2%       1.207µ ± 1%       ~ (p=0.100 n=10)                            1.204µ ± 2%       ~ (p=0.697 n=10)       1.203µ ± 2%       ~ (p=0.898 n=10)
geomean                                                     607.9n            602.3n       -0.91%                                           606.3n       -0.26%                      604.9n       -0.49%

               │ benchmarks/1_fix_preallocations_of_quoted_string.bench │ benchmarks/2_optimize_QuoteBytes.bench │ benchmarks/2_rework_QuoteString_to_iterate_over_bytes.bench │ benchmarks/3_optimize_QuoteBytes.bench │
               │                          B/op                          │     B/op       vs base                 │                B/op                 vs base                 │     B/op       vs base                 │
Sanitize-12                                                  424.0 ± 0%      424.0 ± 0%       ~ (p=1.000 n=10) ¹                           424.0 ± 0%       ~ (p=1.000 n=10) ¹      424.0 ± 0%       ~ (p=1.000 n=10) ¹
SanitizeSQL-12                                               552.0 ± 0%      552.0 ± 0%       ~ (p=1.000 n=10) ¹                           552.0 ± 0%       ~ (p=1.000 n=10) ¹      552.0 ± 0%       ~ (p=1.000 n=10) ¹
geomean                                                      483.8           483.8       +0.00%                                            483.8       +0.00%                       483.8       +0.00%
¹ all samples are equal

               │ benchmarks/1_fix_preallocations_of_quoted_string.bench │ benchmarks/2_optimize_QuoteBytes.bench │ benchmarks/2_rework_QuoteString_to_iterate_over_bytes.bench │ benchmarks/3_optimize_QuoteBytes.bench │
               │                       allocs/op                        │   allocs/op    vs base                 │             allocs/op               vs base                 │   allocs/op    vs base                 │
Sanitize-12                                                  2.000 ± 0%      2.000 ± 0%       ~ (p=1.000 n=10) ¹                           2.000 ± 0%       ~ (p=1.000 n=10) ¹      2.000 ± 0%       ~ (p=1.000 n=10) ¹
SanitizeSQL-12                                               10.00 ± 0%      10.00 ± 0%       ~ (p=1.000 n=10) ¹                           10.00 ± 0%       ~ (p=1.000 n=10) ¹      10.00 ± 0%       ~ (p=1.000 n=10) ¹
geomean                                                      4.472           4.472       +0.00%                                            4.472       +0.00%                       4.472       +0.00%
¹ all samples are equal

sean-

~~[review comment was posted twice for some reason]~~

Additional information warns about using nullable types being used as parameters to query with Valid set to false.

jackc#2120

This removes the import of nanotime via linkname.

When a batch successfully prepared some statements, but then failed to prepare others, the prepared statements that were successfully prepared were not properly cleaned up. This could lead to a "prepared statement already exists" error on subsequent attempts to prepare the same statement. jackc#1847 (comment)

Add the missing 'Z' at the end of the timestamp string, so it can be parsed as timestamp in the RFC3339 format.

make benchmark more extensive add quote to string add BenchmarkSanitizeSQL

fix benchmmark script fix benchmark script

check new quoteBytes

use lexer pool

ninedraft · 2024-12-09T14:31:27Z

@sean I can incorporate your suggestions into this PR if you don't mind

ninedraft marked this pull request as ready for review October 1, 2024 14:52

sean- reviewed Oct 21, 2024

View reviewed changes

mateuszkowalke and others added 20 commits December 9, 2024 16:18

Add additional info for nullable pgtype types

4c1fda0

Additional information warns about using nullable types being used as parameters to query with Valid set to false.

add byte length check to uint32

811b501

Use sql.ErrNoRows as value for pgx.ErrNoRows

1a30a62

Release v5.7.0

19f1994

Fix data race with TraceLog.Config initialization

da51345

jackc#2120

Upgrade puddle to v2.2.2

513a53f

This removes the import of nanotime via linkname.

Update golang.org/x/crypto and golang.org/x/text

be67315

Release v5.7.1

e400c5e

Fix pgtype.Timestamp json unmarshal

fd6496f

Add the missing 'Z' at the end of the timestamp string, so it can be parsed as timestamp in the RFC3339 format.

base case

21392a2

make benchmark more extensive add quote to string add BenchmarkSanitizeSQL

add benchmark tool

d8d0cab

fix benchmmark script fix benchmark script

buf pool

9435a2c

shared bytestring

4f4e892

append AvailableBuffer

39db71a

docs

e142286

quoteBytes

f0180ba

check new quoteBytes

quoteString

c50cb14

decrease number of samples in go benchmark

3a97ffd

add FuzzQuoteString and FuzzQuoteBytes

1ec4baa

ninedraft added 6 commits December 9, 2024 16:18

add lexer and query pools

000ce9c

use lexer pool

rework QuoteString and QuoteBytes as append-style

25a4bd3

add docs to sanitize tests

2f3ae5a

drop too large values from memory pools

9a33a62

add prefix to quoters tests

85d5a4d

fix preallocations of quoted string

97d8358

ninedraft force-pushed the optimize-sanitize branch from 8264272 to 97d8358 Compare December 9, 2024 14:18

optimisations of quote functions by @sean-

174e678

ninedraft force-pushed the optimize-sanitize branch from 50c9eab to 174e678 Compare December 9, 2024 14:34

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Reduce SQL sanitizer allocations #2136

Reduce SQL sanitizer allocations #2136

ninedraft commented Oct 1, 2024

ninedraft commented Oct 1, 2024

jackc commented Oct 5, 2024

vtolstov commented Oct 6, 2024

vtolstov commented Oct 10, 2024

ninedraft commented Oct 15, 2024

jackc commented Oct 18, 2024

sean- left a comment

sean- Oct 21, 2024 •

edited

Loading

sean- Oct 21, 2024

sean- Oct 21, 2024

sean- Oct 21, 2024 •

edited

Loading

sean- left a comment

sean- left a comment •

edited

Loading

ninedraft commented Dec 9, 2024

Reduce SQL sanitizer allocations #2136

Are you sure you want to change the base?

Reduce SQL sanitizer allocations #2136

Conversation

ninedraft commented Oct 1, 2024

ninedraft commented Oct 1, 2024

jackc commented Oct 5, 2024

vtolstov commented Oct 6, 2024

vtolstov commented Oct 10, 2024

ninedraft commented Oct 15, 2024

jackc commented Oct 18, 2024

sean- left a comment

Choose a reason for hiding this comment

sean- Oct 21, 2024 • edited Loading

Choose a reason for hiding this comment

sean- Oct 21, 2024

Choose a reason for hiding this comment

sean- Oct 21, 2024

Choose a reason for hiding this comment

sean- Oct 21, 2024 • edited Loading

Choose a reason for hiding this comment

sean- left a comment

Choose a reason for hiding this comment

sean- left a comment • edited Loading

Choose a reason for hiding this comment

ninedraft commented Dec 9, 2024

sean- Oct 21, 2024 •

edited

Loading

sean- Oct 21, 2024 •

edited

Loading

sean- left a comment •

edited

Loading