From b6f4a3b7019f46bc929f2d9df4887571cb7a8fe4 Mon Sep 17 00:00:00 2001 From: Rafi Shamim Date: Sat, 27 Aug 2022 15:30:20 -0400 Subject: [PATCH] sql/pgwire: fix service latency for COPY Release note (bug fix): Fixed the latency that is reported for COPY comands in the CLI and stats reporting. Release justification: low risk bug fix --- pkg/sql/conn_executor.go | 10 +++++++++ pkg/sql/conn_io.go | 7 ++++++ pkg/sql/pgwire/conn.go | 12 +++++++++- pkg/sql/pgwire/testdata/pgtest/copy | 34 +++++++++++++++++++++++++++++ pkg/testutils/pgtest/datadriven.go | 6 ++++- 5 files changed, 67 insertions(+), 2 deletions(-) diff --git a/pkg/sql/conn_executor.go b/pkg/sql/conn_executor.go index 8d3b6437d2ea..71c61c69d324 100644 --- a/pkg/sql/conn_executor.go +++ b/pkg/sql/conn_executor.go @@ -2047,9 +2047,19 @@ func (ex *connExecutor) execCmd() error { } } case CopyIn: + ex.phaseTimes.SetSessionPhaseTime(sessionphase.SessionQueryReceived, tcmd.TimeReceived) + ex.phaseTimes.SetSessionPhaseTime(sessionphase.SessionStartParse, tcmd.ParseStart) + ex.phaseTimes.SetSessionPhaseTime(sessionphase.SessionEndParse, tcmd.ParseEnd) res = ex.clientComm.CreateCopyInResult(pos) var err error ev, payload, err = ex.execCopyIn(ctx, tcmd) + // Note: we write to ex.statsCollector.phaseTimes, instead of ex.phaseTimes, + // because: + // - stats use ex.statsCollector, not ex.phasetimes. + // - ex.statsCollector merely contains a copy of the times, that + // was created when the statement started executing (via the + // reset() method). + ex.statsCollector.PhaseTimes().SetSessionPhaseTime(sessionphase.SessionQueryServiced, timeutil.Now()) if err != nil { return err } diff --git a/pkg/sql/conn_io.go b/pkg/sql/conn_io.go index e6290165e9f5..40db4caf38ba 100644 --- a/pkg/sql/conn_io.go +++ b/pkg/sql/conn_io.go @@ -328,6 +328,13 @@ type CopyIn struct { // CopyDone is decremented once execution finishes, signaling that control of // the connection is being handed back to the network routine. CopyDone *sync.WaitGroup + // TimeReceived is the time at which the message was received + // from the client. Used to compute the service latency. + TimeReceived time.Time + // ParseStart/ParseEnd are the timing info for parsing of the query. Used for + // stats reporting. + ParseStart time.Time + ParseEnd time.Time } // command implements the Command interface. diff --git a/pkg/sql/pgwire/conn.go b/pkg/sql/pgwire/conn.go index 9b721fa10454..b6ff8d18f85d 100644 --- a/pkg/sql/pgwire/conn.go +++ b/pkg/sql/pgwire/conn.go @@ -860,7 +860,17 @@ func (c *conn) handleSimpleQuery( } copyDone := sync.WaitGroup{} copyDone.Add(1) - if err := c.stmtBuf.Push(ctx, sql.CopyIn{Conn: c, Stmt: cp, CopyDone: ©Done}); err != nil { + if err := c.stmtBuf.Push( + ctx, + sql.CopyIn{ + Conn: c, + Stmt: cp, + CopyDone: ©Done, + TimeReceived: timeReceived, + ParseStart: startParse, + ParseEnd: endParse, + }, + ); err != nil { return err } copyDone.Wait() diff --git a/pkg/sql/pgwire/testdata/pgtest/copy b/pkg/sql/pgwire/testdata/pgtest/copy index 2a811f627a6e..0e4fbdd6fc3c 100644 --- a/pkg/sql/pgwire/testdata/pgtest/copy +++ b/pkg/sql/pgwire/testdata/pgtest/copy @@ -925,3 +925,37 @@ ReadyForQuery {"Type":"DataRow","Values":[{"text":"{2,3,4,5,6}"}]} {"Type":"CommandComplete","CommandTag":"SELECT 2"} {"Type":"ReadyForQuery","TxStatus":"I"} + + +send crdb_only +Query {"String": "COPY c(d) FROM STDIN WITH CSV"} +CopyData {"Data": "\"{0,1}\"\n"} +CopyData {"Data": "\"{2,3,4,5,6}\"\n"} +CopyData {"Data": "\\.\n"} +CopyDone +---- + + +until crdb_only ignore=RowDescription +ReadyForQuery +---- +{"Type":"CopyInResponse","ColumnFormatCodes":[0]} +{"Type":"CommandComplete","CommandTag":"COPY 2"} +{"Type":"ReadyForQuery","TxStatus":"I"} + +# SHOW LAST QUERY STATISTICS cannot be used as a statement source, so we +# store the result as a variable in the test. +let $copy_service_latency crdb_only +Query {"String": "SHOW LAST QUERY STATISTICS RETURNING service_latency"} +---- + +send crdb_only +Query {"String": "SELECT '$copy_service_latency'::INTERVAL > '0 seconds'::INTERVAL"} +---- + +until crdb_only ignore=RowDescription +ReadyForQuery +---- +{"Type":"DataRow","Values":[{"text":"t"}]} +{"Type":"CommandComplete","CommandTag":"SELECT 1"} +{"Type":"ReadyForQuery","TxStatus":"I"} diff --git a/pkg/testutils/pgtest/datadriven.go b/pkg/testutils/pgtest/datadriven.go index 212769b7838c..73767b64d3c7 100644 --- a/pkg/testutils/pgtest/datadriven.go +++ b/pkg/testutils/pgtest/datadriven.go @@ -92,7 +92,11 @@ func RunTest(t *testing.T, path, addr, user string) { return d.Expected case "let": - require.Len(t, d.CmdArgs, 1, "only one argument permitted for let") + if (d.HasArg("crdb_only") && !p.isCockroachDB) || + (d.HasArg("noncrdb_only") && p.isCockroachDB) { + return d.Expected + } + require.GreaterOrEqual(t, len(d.CmdArgs), 1, "at least one argument required for let") require.Truef(t, strings.HasPrefix(d.CmdArgs[0].Key, "$"), "let argument must begin with '$'") lines := strings.Split(d.Input, "\n") require.Len(t, lines, 1, "only one input command permitted for let")