Skip to content

Commit b5971c6

Browse files
asmyasnikovasmyasnikov
authored
Merge pull request #84 from ydb-platform/stream-result
drop creation of goroutine on each stream call - stream.Recv() called…
2 parents 685a635 + 0b24cd2 commit b5971c6

File tree

8 files changed

+100
-147
lines changed

8 files changed

+100
-147
lines changed

‎CHANGELOG.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
* Added `retry.RetryableError()` for returns user-defined error which must be retryed
1212
* Renamed internal type `internal.errors.OperationCompleted` to `internal.errors.OperationStatus`
1313
* Added `String()` method to `table.KeyRange` and `table.Value` types
14+
* Replaced creation of goroutine on each stream call to explicit call stream.Recv() on NextResultSet()
1415

1516
## 3.6.2
1617
* Refactored table retry helpers

‎internal/conn/conn.go‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -335,12 +335,12 @@ func (c *conn) NewStream(
335335

336336
var s grpc.ClientStream
337337
s, err = cc.NewStream(ctx, desc, method, append(opts, grpc.MaxCallRecvMsgSize(50*1024*1024))...)
338-
c.release(ctx)
339338
if err != nil {
340339
err = errors.MapGRPCError(err)
341340
if errors.MustPessimizeEndpoint(err) {
342341
c.pessimize(ctx, err)
343342
}
343+
c.release(ctx)
344344
return nil, err
345345
}
346346

@@ -349,6 +349,7 @@ func (c *conn) NewStream(
349349
s: s,
350350
onDone: func(ctx context.Context) {
351351
cancel()
352+
c.release(ctx)
352353
},
353354
recv: streamRecv,
354355
}, nil

‎internal/scripting/scripting.go‎

Lines changed: 20 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -2,15 +2,15 @@ package scheme
22

33
import (
44
"context"
5-
"io"
65

76
"google.golang.org/grpc"
87
"google.golang.org/protobuf/proto"
98

109
"github.com/ydb-platform/ydb-go-genproto/Ydb_Scripting_V1"
10+
"github.com/ydb-platform/ydb-go-genproto/protos/Ydb"
1111
"github.com/ydb-platform/ydb-go-genproto/protos/Ydb_Scripting"
12+
"github.com/ydb-platform/ydb-go-genproto/protos/Ydb_TableStats"
1213

13-
"github.com/ydb-platform/ydb-go-sdk/v3/internal/errors"
1414
"github.com/ydb-platform/ydb-go-sdk/v3/internal/table/scanner"
1515
"github.com/ydb-platform/ydb-go-sdk/v3/internal/value"
1616
"github.com/ydb-platform/ydb-go-sdk/v3/scripting"
@@ -107,43 +107,29 @@ func (c *client) StreamExecute(
107107
return nil, err
108108
}
109109

110-
r := scanner.NewStream()
111-
go func() {
112-
var (
113-
response *Ydb_Scripting.ExecuteYqlPartialResponse
114-
err error
115-
)
116-
defer func() {
117-
cancel()
118-
r.Close()
119-
}()
120-
for {
110+
return scanner.NewStream(
111+
func(ctx context.Context) (
112+
set *Ydb.ResultSet,
113+
stats *Ydb_TableStats.QueryStats,
114+
err error,
115+
) {
121116
select {
122117
case <-ctx.Done():
123-
err = ctx.Err()
124-
r.SetErr(err)
125-
return
118+
return nil, nil, ctx.Err()
126119
default:
127-
if response, err = stream.Recv(); err != nil {
128-
if !errors.Is(err, io.EOF) {
129-
r.SetErr(err)
130-
// nolint:ineffassign
131-
err = nil
132-
}
133-
return
134-
}
135-
if result := response.GetResult(); result != nil {
136-
if resultSet := result.GetResultSet(); resultSet != nil {
137-
r.Append(resultSet)
138-
}
139-
if stats := result.GetQueryStats(); stats != nil {
140-
r.UpdateStats(stats)
141-
}
120+
response, err := stream.Recv()
121+
result := response.GetResult()
122+
if result == nil || err != nil {
123+
return nil, nil, err
142124
}
125+
return result.GetResultSet(), result.GetQueryStats(), nil
143126
}
144-
}
145-
}()
146-
return r, nil
127+
},
128+
func(err error) error {
129+
cancel()
130+
return err
131+
},
132+
), nil
147133
}
148134

149135
func (c *client) Close(context.Context) error {

‎internal/table/scanner/result.go‎

Lines changed: 28 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -2,20 +2,23 @@ package scanner
22

33
import (
44
"context"
5+
"io"
56
"sync"
67
"time"
78

89
"github.com/ydb-platform/ydb-go-genproto/protos/Ydb"
910
"github.com/ydb-platform/ydb-go-genproto/protos/Ydb_TableStats"
1011

12+
"github.com/ydb-platform/ydb-go-sdk/v3/internal/errors"
1113
public "github.com/ydb-platform/ydb-go-sdk/v3/table/result"
1214
"github.com/ydb-platform/ydb-go-sdk/v3/table/stats"
1315
)
1416

1517
type result struct {
1618
scanner
1719

18-
stats *Ydb_TableStats.QueryStats
20+
statsMtx sync.RWMutex
21+
stats *Ydb_TableStats.QueryStats
1922

2023
closedMtx sync.RWMutex
2124
closed bool
@@ -24,7 +27,8 @@ type result struct {
2427
type streamResult struct {
2528
result
2629

27-
ch chan *Ydb.ResultSet
30+
recv func(ctx context.Context) (*Ydb.ResultSet, *Ydb_TableStats.QueryStats, error)
31+
close func(error) error
2832
}
2933

3034
type unaryResult struct {
@@ -55,19 +59,6 @@ func (r *result) isClosed() bool {
5559
return r.closed
5660
}
5761

58-
func (r *result) UpdateStats(stats *Ydb_TableStats.QueryStats) {
59-
r.stats = stats
60-
}
61-
62-
func (r *streamResult) Append(set *Ydb.ResultSet) {
63-
r.closedMtx.RLock()
64-
defer r.closedMtx.RUnlock()
65-
if r.closed {
66-
return
67-
}
68-
r.ch <- set
69-
}
70-
7162
type resultWithError interface {
7263
SetErr(err error)
7364
}
@@ -80,14 +71,15 @@ type UnaryResult interface {
8071
type StreamResult interface {
8172
public.StreamResult
8273
resultWithError
83-
84-
Append(set *Ydb.ResultSet)
85-
UpdateStats(stats *Ydb_TableStats.QueryStats)
8674
}
8775

88-
func NewStream() StreamResult {
76+
func NewStream(
77+
recv func(ctx context.Context) (*Ydb.ResultSet, *Ydb_TableStats.QueryStats, error),
78+
onClose func(error) error,
79+
) StreamResult {
8980
r := &streamResult{
90-
ch: make(chan *Ydb.ResultSet, 1),
81+
recv: recv,
82+
close: onClose,
9183
}
9284
return r
9385
}
@@ -128,15 +120,11 @@ func (r *streamResult) NextResultSet(ctx context.Context, columns ...string) boo
128120
if r.inactive() {
129121
return false
130122
}
131-
select {
132-
case s, ok := <-r.ch:
133-
if !ok || s == nil {
134-
return false
135-
}
136-
r.Reset(s, columns...)
137-
return true
138-
139-
case <-ctx.Done():
123+
s, stats, err := r.recv(ctx)
124+
if errors.Is(err, io.EOF) {
125+
return false
126+
}
127+
if err != nil {
140128
r.errMtx.Lock()
141129
if r.err == nil {
142130
r.err = ctx.Err()
@@ -145,6 +133,13 @@ func (r *streamResult) NextResultSet(ctx context.Context, columns ...string) boo
145133
r.Reset(nil)
146134
return false
147135
}
136+
r.Reset(s, columns...)
137+
if stats != nil {
138+
r.statsMtx.Lock()
139+
r.stats = stats
140+
r.statsMtx.Unlock()
141+
}
142+
return true
148143
}
149144

150145
// CurrentResultSet get current result set
@@ -155,8 +150,10 @@ func (r *result) CurrentResultSet() public.Set {
155150
// Stats returns query execution queryStats.
156151
func (r *result) Stats() stats.QueryStats {
157152
var s queryStats
153+
r.statsMtx.RLock()
158154
s.stats = r.stats
159-
s.processCPUTime = time.Microsecond * time.Duration(r.stats.GetProcessCpuTimeUs())
155+
r.statsMtx.RUnlock()
156+
s.processCPUTime = time.Microsecond * time.Duration(s.stats.GetProcessCpuTimeUs())
160157
s.pos = 0
161158
return &s
162159
}
@@ -169,8 +166,7 @@ func (r *streamResult) Close() (err error) {
169166
return nil
170167
}
171168
r.closed = true
172-
close(r.ch)
173-
return nil
169+
return r.close(r.Err())
174170
}
175171

176172
func (r *result) inactive() bool {

‎internal/table/session.go‎

Lines changed: 41 additions & 63 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,6 @@ package table
22

33
import (
44
"context"
5-
"io"
65
"net/url"
76
"strconv"
87
"sync"
@@ -14,6 +13,7 @@ import (
1413
"github.com/ydb-platform/ydb-go-genproto/Ydb_Table_V1"
1514
"github.com/ydb-platform/ydb-go-genproto/protos/Ydb"
1615
"github.com/ydb-platform/ydb-go-genproto/protos/Ydb_Table"
16+
"github.com/ydb-platform/ydb-go-genproto/protos/Ydb_TableStats"
1717

1818
"github.com/ydb-platform/ydb-go-sdk/v3/internal/cluster"
1919
"github.com/ydb-platform/ydb-go-sdk/v3/internal/errors"
@@ -700,43 +700,34 @@ func (s *session) StreamReadTable(
700700
onDone := trace.TableOnSessionQueryStreamRead(s.trace, &ctx, s)
701701
if err != nil {
702702
cancel()
703-
onDone(nil, err)
703+
onDone(err)
704704
return nil, err
705705
}
706706

707-
r := scanner.NewStream()
708-
go func() {
709-
var (
710-
response *Ydb_Table.ReadTableResponse
711-
err error
712-
)
713-
defer func() {
714-
cancel()
715-
onDone(r, errors.HideEOF(err))
716-
r.Close()
717-
}()
718-
for {
707+
return scanner.NewStream(
708+
func(ctx context.Context) (
709+
set *Ydb.ResultSet,
710+
stats *Ydb_TableStats.QueryStats,
711+
err error,
712+
) {
719713
select {
720714
case <-ctx.Done():
721-
err = ctx.Err()
722-
r.SetErr(err)
723-
return
715+
return nil, nil, ctx.Err()
724716
default:
725-
if response, err = stream.Recv(); err != nil {
726-
if !errors.Is(err, io.EOF) {
727-
r.SetErr(err)
728-
}
729-
return
730-
}
731-
if result := response.GetResult(); result != nil {
732-
if resultSet := result.GetResultSet(); resultSet != nil {
733-
r.Append(resultSet)
734-
}
717+
response, err := stream.Recv()
718+
result := response.GetResult()
719+
if result == nil || err != nil {
720+
return nil, nil, err
735721
}
722+
return result.GetResultSet(), nil, nil
736723
}
737-
}
738-
}()
739-
return r, nil
724+
},
725+
func(err error) error {
726+
cancel()
727+
onDone(err)
728+
return err
729+
},
730+
), nil
740731
}
741732

742733
// StreamExecuteScanQuery scan-reads table at given path with given options.
@@ -771,47 +762,34 @@ func (s *session) StreamExecuteScanQuery(
771762
onDone := trace.TableOnSessionQueryStreamExecute(s.trace, &ctx, s, q, params)
772763
if err != nil {
773764
cancel()
774-
onDone(nil, err)
765+
onDone(err)
775766
return nil, err
776767
}
777768

778-
r := scanner.NewStream()
779-
go func() {
780-
var (
781-
response *Ydb_Table.ExecuteScanQueryPartialResponse
782-
err error
783-
)
784-
defer func() {
785-
cancel()
786-
onDone(r, errors.HideEOF(err))
787-
r.Close()
788-
}()
789-
for {
769+
return scanner.NewStream(
770+
func(ctx context.Context) (
771+
set *Ydb.ResultSet,
772+
stats *Ydb_TableStats.QueryStats,
773+
err error,
774+
) {
790775
select {
791776
case <-ctx.Done():
792-
err = ctx.Err()
793-
r.SetErr(err)
794-
return
777+
return nil, nil, ctx.Err()
795778
default:
796-
if response, err = stream.Recv(); err != nil {
797-
if !errors.Is(err, io.EOF) {
798-
r.SetErr(err)
799-
err = nil
800-
}
801-
return
802-
}
803-
if result := response.GetResult(); result != nil {
804-
if resultSet := result.GetResultSet(); resultSet != nil {
805-
r.Append(resultSet)
806-
}
807-
if stats := result.GetQueryStats(); stats != nil {
808-
r.UpdateStats(stats)
809-
}
779+
response, err := stream.Recv()
780+
result := response.GetResult()
781+
if result == nil || err != nil {
782+
return nil, nil, err
810783
}
784+
return result.GetResultSet(), result.GetQueryStats(), nil
811785
}
812-
}
813-
}()
814-
return r, nil
786+
},
787+
func(err error) error {
788+
cancel()
789+
onDone(err)
790+
return err
791+
},
792+
), nil
815793
}
816794

817795
// BulkUpsert uploads given list of ydb struct values to the table.

0 commit comments

Comments
 (0)