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]~~

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

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

vtolstov · 2024-12-25T23:16:17Z

may be we can move forward with this pr?

lrweck · 2024-12-27T12:17:46Z

conn_test.go

+	require.Equal(t, "no rows in result set", pgx.ErrNoRows.Error())
+
+	require.ErrorIs(t, pgx.ErrNoRows, sql.ErrNoRows, "pgx.ErrNowRows must match sql.ErrNoRows")
+	require.ErrorIs(t, pgx.ErrNoRows, pgx.ErrNoRows, "sql.ErrNowRows must match pgx.ErrNoRows")


Suggested change

require.ErrorIs(t, pgx.ErrNoRows, pgx.ErrNoRows, "sql.ErrNowRows must match pgx.ErrNoRows")

require.ErrorIs(t, sql.ErrNoRows, pgx.ErrNoRows, "sql.ErrNowRows must match pgx.ErrNoRows")

Thanks, but it's an erroneous test case, which I forgot to delete (I'm deleting it now).
The only thing we must guarantee is

errors.Is(pgx.ErrNoRows, sql.ErrNoRows) == true

errors.Is(sql.ErrNoRows, pgx.ErrNoRows) == true is achievable, but is not required. I rather not at add it since it will require some extra error-matching magic

lrweck · 2024-12-27T12:22:03Z

pgtype/timestamp.go

@@ -104,8 +104,8 @@ func (ts *Timestamp) UnmarshalJSON(b []byte) error {
 	case "-infinity":
 		*ts = Timestamp{Valid: true, InfinityModifier: -Infinity}
 	default:
-		// PostgreSQL uses ISO 8601 for to_json function and casting from a string to timestamptz
-		tim, err := time.Parse(time.RFC3339Nano, *s)
+		// PostgreSQL uses ISO 8601 wihout timezone for to_json function and casting from a string to timestampt


Suggested change

// PostgreSQL uses ISO 8601 wihout timezone for to_json function and casting from a string to timestampt

// PostgreSQL uses ISO 8601 without timezone for to_json function and casting from a string to timestamptz

Thanks! I thing I'll extract it as a separate PR - want to keep scope of this one narrow

make benchmark more extensive add quote to string add BenchmarkSanitizeSQL

fix benchmmark script fix benchmark script

check new quoteBytes

use lexer pool

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

sean- reviewed Oct 21, 2024

View reviewed changes

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

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

lrweck reviewed Dec 27, 2024

View reviewed changes

ninedraft added 15 commits December 28, 2024 13:31

base case

afa974f

make benchmark more extensive add quote to string add BenchmarkSanitizeSQL

add benchmark tool

aabed18

fix benchmmark script fix benchmark script

buf pool

efc2c9f

shared bytestring

546ad2f

append AvailableBuffer

ee718a1

docs

1752f7b

quoteBytes

58d4c0c

check new quoteBytes

quoteString

ea1e13a

decrease number of samples in go benchmark

4293b25

add FuzzQuoteString and FuzzQuoteBytes

c4c1076

add lexer and query pools

39ffc8b

use lexer pool

rework QuoteString and QuoteBytes as append-style

59d6aa8

add docs to sanitize tests

90a77b1

drop too large values from memory pools

47cbd8e

add prefix to quoters tests

057937d

ninedraft added 2 commits December 28, 2024 13:31

fix preallocations of quoted string

120c89f

optimisations of quote functions by @sean-

da0315d

ninedraft force-pushed the optimize-sanitize branch from 174e678 to da0315d Compare December 28, 2024 11:31

TestErrNoRows: remove bad test case

e452f80

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

ninedraft Dec 28, 2024

sean- Oct 21, 2024

ninedraft Dec 28, 2024

sean- Oct 21, 2024

ninedraft Dec 28, 2024

sean- Oct 21, 2024 •

edited

Loading

ninedraft Dec 28, 2024

sean- left a comment

sean- left a comment •

edited

Loading

ninedraft commented Dec 9, 2024

vtolstov commented Dec 25, 2024

lrweck Dec 27, 2024

ninedraft Dec 28, 2024

lrweck Dec 27, 2024

ninedraft Dec 28, 2024

	require.ErrorIs(t, pgx.ErrNoRows, pgx.ErrNoRows, "sql.ErrNowRows must match pgx.ErrNoRows")
	require.ErrorIs(t, sql.ErrNoRows, pgx.ErrNoRows, "sql.ErrNowRows must match pgx.ErrNoRows")

	// PostgreSQL uses ISO 8601 wihout timezone for to_json function and casting from a string to timestampt
	// PostgreSQL uses ISO 8601 without timezone for to_json function and casting from a string to timestamptz

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

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

sean- Oct 21, 2024 • edited Loading

Choose a reason for hiding this comment

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

vtolstov commented Dec 25, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

sean- Oct 21, 2024 •

edited

Loading

sean- Oct 21, 2024 •

edited

Loading

sean- left a comment •

edited

Loading