From ea1ab93882d480955e7466eafbb80875755ff4c5 Mon Sep 17 00:00:00 2001 From: Ganeshprasad Biradar Date: Tue, 29 Jun 2021 19:03:33 +0530 Subject: [PATCH] pgwire: pre-allocate a temporary buffer to use when serializing arrays/tuples A new "writeBuffer" was being allocated every time when serializing an array/tuple. This is very heavy and unecessary. To address this, this patch creates a temporary buffer in the existing "writeBuffer" struct which is used specially for serializing array/tuple. This is more efficient than allocating a new temporary "writebuffer" every time. Release note (performance improvement): A special temporary buffer was created to be used only for serializing arrays/tuples. This provided better performance. The increase in performance was seen by running the BenchmarkEncodings benchmark for arrays and tuples. Fixes: #21711 --- pkg/sql/pgwire/types.go | 124 ++++++++++++++++++++++----------- pkg/sql/pgwire/write_buffer.go | 11 ++- 2 files changed, 95 insertions(+), 40 deletions(-) diff --git a/pkg/sql/pgwire/types.go b/pkg/sql/pgwire/types.go index e098de65425c..bde7d3b97738 100644 --- a/pkg/sql/pgwire/types.go +++ b/pkg/sql/pgwire/types.go @@ -130,6 +130,89 @@ func writeTextTimestampTZ(b *writeBuffer, v time.Time, sessionLoc *time.Location b.write(s) } +func writeTuple( + ctx context.Context, b *writeBuffer, v *tree.DTuple, sessionLoc *time.Location, t *types.T, +) { + // initialState will be used to decide whether to swap the wrapped and the scratch buffer. + initialState := b.usingScratch + if b.usingScratch == false { + temp := b.wrapped + b.wrapped = b.scratch + b.scratch = temp + b.usingScratch = true + } + + // Put the number of datums. + b.putInt32(int32(len(v.D))) + tupleTypes := t.TupleContents() + for i, elem := range v.D { + oid := tupleTypes[i].Oid() + b.putInt32(int32(oid)) + b.writeBinaryDatum(ctx, elem, sessionLoc, tupleTypes[i]) + } + + if initialState == false { + + // We are done with recursive calls, copy scratch buffer to wrapped buffer and reset it. + temp := b.wrapped + b.wrapped = b.scratch + b.scratch = temp + b.writeLengthPrefixedBuffer(&b.scratch) + b.scratch.Reset() + b.usingScratch = false + } +} + +func writeArray( + ctx context.Context, b *writeBuffer, v *tree.DArray, sessionLoc *time.Location, t *types.T, +) { + if v.ParamTyp.Family() == types.ArrayFamily { + b.setError(unimplemented.NewWithIssueDetail(32552, + "binenc", "unsupported binary serialization of multidimensional arrays")) + return + } + + // Look at writeTuple for more info on initialState. + initialState := b.usingScratch + if b.usingScratch == false { + temp := b.wrapped + b.wrapped = b.scratch + b.scratch = temp + b.usingScratch = true + } + + // Put the number of dimensions. We currently support 1d arrays only. + var ndims int32 = 1 + if v.Len() == 0 { + ndims = 0 + } + b.putInt32(ndims) + hasNulls := 0 + if v.HasNulls { + hasNulls = 1 + } + oid := v.ParamTyp.Oid() + b.putInt32(int32(hasNulls)) + b.putInt32(int32(oid)) + if v.Len() > 0 { + b.putInt32(int32(v.Len())) + // Lower bound, we only support a lower bound of 1. + b.putInt32(1) + for _, elem := range v.Array { + b.writeBinaryDatum(ctx, elem, sessionLoc, v.ParamTyp) + } + } + + if initialState == false { + temp := b.wrapped + b.wrapped = b.scratch + b.scratch = temp + b.writeLengthPrefixedBuffer(&b.scratch) + b.scratch.Reset() + b.usingScratch = false + } +} + // writeTextDatum writes d to the buffer. Type t must be specified for types // that have various width encodings and therefore need padding (chars). // It is ignored (and can be nil) for types which do not need padding. @@ -649,16 +732,7 @@ func writeBinaryDatumNotNull( case *tree.DTuple: // TODO(andrei): We shouldn't be allocating a new buffer for every array. - subWriter := newWriteBuffer(nil /* bytecount */) - // Put the number of datums. - subWriter.putInt32(int32(len(v.D))) - tupleTypes := t.TupleContents() - for i, elem := range v.D { - oid := tupleTypes[i].Oid() - subWriter.putInt32(int32(oid)) - subWriter.writeBinaryDatum(ctx, elem, sessionLoc, tupleTypes[i]) - } - b.writeLengthPrefixedBuffer(&subWriter.wrapped) + writeTuple(ctx, b, v, sessionLoc, t) case *tree.DBox2D: b.putInt32(32) @@ -676,35 +750,7 @@ func writeBinaryDatumNotNull( b.write(v.EWKB()) case *tree.DArray: - if v.ParamTyp.Family() == types.ArrayFamily { - b.setError(unimplemented.NewWithIssueDetail(32552, - "binenc", "unsupported binary serialization of multidimensional arrays")) - return - } - // TODO(andrei): We shouldn't be allocating a new buffer for every array. - subWriter := newWriteBuffer(nil /* bytecount */) - // Put the number of dimensions. We currently support 1d arrays only. - var ndims int32 = 1 - if v.Len() == 0 { - ndims = 0 - } - subWriter.putInt32(ndims) - hasNulls := 0 - if v.HasNulls { - hasNulls = 1 - } - oid := v.ParamTyp.Oid() - subWriter.putInt32(int32(hasNulls)) - subWriter.putInt32(int32(oid)) - if v.Len() > 0 { - subWriter.putInt32(int32(v.Len())) - // Lower bound, we only support a lower bound of 1. - subWriter.putInt32(1) - for _, elem := range v.Array { - subWriter.writeBinaryDatum(ctx, elem, sessionLoc, v.ParamTyp) - } - } - b.writeLengthPrefixedBuffer(&subWriter.wrapped) + writeArray(ctx, b, v, sessionLoc, t) case *tree.DJSON: writeBinaryJSON(b, v.JSON) diff --git a/pkg/sql/pgwire/write_buffer.go b/pkg/sql/pgwire/write_buffer.go index 695429619496..0dcd1790694c 100644 --- a/pkg/sql/pgwire/write_buffer.go +++ b/pkg/sql/pgwire/write_buffer.go @@ -28,7 +28,15 @@ type writeBuffer struct { _ util.NoCopy wrapped bytes.Buffer - err error + + // scratch will be used to store temporary data when serializing array/tuple. + scratch bytes.Buffer + + // usingScratch will be used to keep track of which buffer we are currently using + // i.e either wrapped or scratch. + usingScratch bool + + err error // Buffer used for temporary storage. putbuf [64]byte @@ -53,6 +61,7 @@ func (b *writeBuffer) init(bytecount *metric.Counter) { b.bytecount = bytecount b.textFormatter = tree.NewFmtCtx(tree.FmtPgwireText) b.simpleFormatter = tree.NewFmtCtx(tree.FmtSimple) + b.usingScratch = false } // Write implements the io.Write interface.