Merge #147861 #148222

147861: storage: introduce ComputeSSTStatsDiff r=sumeerbhola a=msbutler

This patch adds the new ComputeSSTStatsDiff utility which will be called in batcheval.EvalAddSSTable in a future PR to ingest an sst over non empty key space with an accurate stats update.

Informs #145548

Release note: none

148222: pgwire: limit size of SSL frames sent r=fqazi a=fqazi

Previously, CRDB bumped up the result buffer size
(sql.defaults.results_buffer.size) from 16k to a higher value. After this change we started seeing asynchronous libpq clients start hanging because they can only consume 8 kilobytes of data at a time. As a result large SSL frames can be problematic in this configuration. To address this, this patch limits CRDB to sending the result buffer at 8 kilobytes at a time.

Fixes: #147890

Release note (bug fix): Addressed a bug where libpq clients using the async API could hang with large result sets (psycopg, activerecord, ruby-pg)

Co-authored-by: Michael Butler <[email protected]>
Co-authored-by: Faizan Qazi <[email protected]>

Author: 3 people

pkg/acceptance/adapter_test.go

@@ -69,6 +69,7 @@ func TestDockerNodeJS(t *testing.T) {
 	export SHOULD_FAIL=%v
 	# Get access to globally installed node modules.
 	export NODE_PATH=$NODE_PATH:/usr/lib/node
+	export NODE_TLS_REJECT_UNAUTHORIZED=0
 	# Have a 10 second timeout on promises, in case the server is slow.
 	/usr/lib/node/.bin/mocha -t 10000 . 
 	`
@@ -110,9 +111,12 @@ func TestDockerPython(t *testing.T) {
 	ctx := context.Background()
 	t.Run("Success", func(t *testing.T) {
 		testDockerSuccess(ctx, t, "python", []string{"sh", "-c", "cd /mnt/data/python && python test.py 3"})
+		testDockerSuccess(ctx, t, "python", []string{"sh", "-c", "cd /mnt/data/python && python3 test_pyscopg3.py 3"})
+
 	})
 	t.Run("Fail", func(t *testing.T) {
 		testDockerFail(ctx, t, "python", []string{"sh", "-c", "cd /mnt/data/python && python test.py 2"})
+		testDockerFail(ctx, t, "python", []string{"sh", "-c", "cd /mnt/data/python && python3 test_pyscopg3.py 2"})
 	})
 }
 

pkg/acceptance/testdata/Dockerfile

@@ -63,7 +63,13 @@ RUN curl -fsSL https://dl.yarnpkg.com/debian/pubkey.gpg > /etc/apt/trusted.gpg.d
     ruby \
     ruby-pg \
     xmlstarlet \
-    yarn
+    yarn \
+    python3-dev \
+    python3-pip \
+    python3-setuptools \
+&& pip3 install --upgrade pip \
+&& pip3 install backports.zoneinfo \
+&& pip3 install psycopg
 
 RUN case ${TARGETPLATFORM} in \
     "linux/amd64") ARCH=amd64; SHASUM=442dae58b727a79f81368127fac141d7f95501ffa05f8c48943d27c4e807deb7 ;; \

pkg/acceptance/testdata/node/base-test.js

@@ -71,6 +71,10 @@ describe('arrays', () => {
   });
 });
 
+// Temporarily disabled until https://github.com/brianc/node-postgres/issues/3487
+// gets resolved. The binary encoding in node-postgres was regressed at some
+// point leading to protocol violation errors.
+/*
 describe('regression tests', () => {
   it('allows you to switch between format modes for arrays', () => {
     return client.query({
@@ -88,3 +92,4 @@ describe('regression tests', () => {
           });
   });
 })
+*/

pkg/acceptance/testdata/python/test_pyscopg3.py

@@ -0,0 +1,71 @@
+# Copyright 2018 The Cockroach Authors.
+#
+# Use of this software is governed by the CockroachDB Software License
+# included in the /LICENSE file.
+
+import decimal
+import sys
+import psycopg
+
+conn = psycopg.connect('')
+cur = conn.cursor()
+cur.execute("SELECT 1, 2+{}".format(sys.argv[1]))
+v = cur.fetchall()
+assert v == [(1, 5)]
+
+# Verify #6597 (timestamp format) is fixed.
+cur = conn.cursor()
+cur.execute("SELECT now()")
+v = cur.fetchall()
+
+# Verify round-trip of strings containing backslashes.
+# https://github.com/cockroachdb/cockroachdb-python/issues/23
+s = ('\\\\',)
+cur.execute("SELECT %s::STRING", s)
+v = cur.fetchall()
+assert v == [s], (v, s)
+
+# Verify decimals with exponents can be parsed.
+cur = conn.cursor()
+cur.execute("SELECT 1e1::decimal")
+v = cur.fetchall()
+d = v[0][0]
+assert type(d) is decimal.Decimal
+# Use of compare_total here guarantees that we didn't just get '10' back, we got '1e1'.
+assert d.compare_total(decimal.Decimal('1e1')) == 0
+
+# Verify arrays with strings can be parsed.
+cur = conn.cursor()
+cur.execute("SELECT ARRAY['foo','bar','baz']")
+v = cur.fetchall()
+d = v[0][0]
+assert d == ["foo","bar","baz"]
+
+# Verify JSON values come through properly.
+cur = conn.cursor()
+cur.execute("SELECT '{\"a\":\"b\"}'::JSONB")
+v = cur.fetchall()
+d = v[0][0]
+assert d == {"a": "b"}
+
+
+# Verify queries with large result sets can be run and do not
+# hang.
+query = """
+    SELECT
+        '000000000000000000000' AS aaaaaaaa,
+        '000000000000000000000' AS b,
+        '000000000000000000000' AS c,
+        '000000000000000000000' AS d,
+        '000000000000000000000' AS e,
+        '000000000000000000000' AS f,
+        '000000000000000000000' AS g,
+        '0000000000000000000' AS h
+    FROM generate_series(1, 238)
+"""
+
+for i in range(30):
+  cur.execute(query)
+  v = cur.fetchall()
+  assert len(v) == 238
+

pkg/acceptance/util_docker.go

@@ -63,7 +63,7 @@ func testDockerSuccess(ctx context.Context, t *testing.T, name string, cmd []str
 const (
 	// Iterating against a locally built version of the docker image can be done
 	// by changing acceptanceImage to the hash of the container.
-	acceptanceImage = "us-east1-docker.pkg.dev/crl-ci-images/cockroach/acceptance:20221005-223354"
+	acceptanceImage = "us-east1-docker.pkg.dev/crl-ci-images/cockroach/acceptance:20250612-132728"
 )
 
 func testDocker(

pkg/sql/pgwire/conn.go

@@ -1251,11 +1251,26 @@ func (c *conn) Flush(pos sql.CmdPos) error {
 	// the underlying ring buffer memory for reuse.
 	c.writerState.fi.cmdStarts.Reset()
 
-	_ /* n */, err := c.writerState.buf.WriteTo(c.conn)
-	if err != nil {
-		c.setErr(err)
-		return err
+	// When flushing, writes into the socket buffer limit how many bytes
+	// are written into an SSL frame at a time. Stock Postgres will at most
+	// write 8 kilobytes of unencrypted data at a time. libpq can also only
+	// consume 8 kilobytes (via PQConsumeInput), so async clients will hang
+	// if bigger SSL frames are sent.
+	const maxWriteSize = 8192
+	bytesToWrite := c.writerState.buf.Bytes()
+	for i := 0; i < len(bytesToWrite); i += maxWriteSize {
+		targetSize := min(len(bytesToWrite)-i, maxWriteSize)
+		m, err := c.conn.Write(bytesToWrite[i : i+targetSize])
+		if err == nil && m < targetSize {
+			err = io.ErrShortWrite
+		}
+		if err != nil {
+			c.setErr(err)
+			c.writerState.buf.Reset()
+			return err
+		}
 	}
+	c.writerState.buf.Reset()
 	return nil
 }
 

pkg/storage/BUILD.bazel

@@ -148,6 +148,7 @@ go_test(
         "pebble_mvcc_scanner_test.go",
         "pebble_test.go",
         "read_as_of_iterator_test.go",
+        "sst_stats_diff_test.go",
         "sst_test.go",
         "sst_writer_test.go",
         "store_properties_test.go",