vitess.io/vitess@v0.16.2/go/vt/vttablet/tabletserver/vstreamer/resultstreamer.go (about) 1 /* 2 Copyright 2019 The Vitess Authors. 3 4 Licensed under the Apache License, Version 2.0 (the "License"); 5 you may not use this file except in compliance with the License. 6 You may obtain a copy of the License at 7 8 http://www.apache.org/licenses/LICENSE-2.0 9 10 Unless required by applicable law or agreed to in writing, software 11 distributed under the License is distributed on an "AS IS" BASIS, 12 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 See the License for the specific language governing permissions and 14 limitations under the License. 15 */ 16 17 package vstreamer 18 19 import ( 20 "context" 21 "fmt" 22 "time" 23 24 "vitess.io/vitess/go/sqltypes" 25 "vitess.io/vitess/go/vt/dbconfigs" 26 binlogdatapb "vitess.io/vitess/go/vt/proto/binlogdata" 27 "vitess.io/vitess/go/vt/sqlparser" 28 ) 29 30 // resultStreamer streams the results of the requested query 31 // along with the GTID of the snapshot. This is used by vdiff 32 // to synchronize the target to that GTID before comparing 33 // the results. 34 type resultStreamer struct { 35 ctx context.Context 36 cancel func() 37 38 cp dbconfigs.Connector 39 query string 40 tableName sqlparser.IdentifierCS 41 send func(*binlogdatapb.VStreamResultsResponse) error 42 vse *Engine 43 pktsize PacketSizer 44 } 45 46 func newResultStreamer(ctx context.Context, cp dbconfigs.Connector, query string, send func(*binlogdatapb.VStreamResultsResponse) error, vse *Engine) *resultStreamer { 47 ctx, cancel := context.WithCancel(ctx) 48 return &resultStreamer{ 49 ctx: ctx, 50 cancel: cancel, 51 cp: cp, 52 query: query, 53 send: send, 54 vse: vse, 55 pktsize: DefaultPacketSizer(), 56 } 57 } 58 59 func (rs *resultStreamer) Cancel() { 60 rs.cancel() 61 } 62 63 func (rs *resultStreamer) Stream() error { 64 _, fromTable, err := analyzeSelect(rs.query) 65 if err != nil { 66 return err 67 } 68 rs.tableName = fromTable 69 70 conn, err := snapshotConnect(rs.ctx, rs.cp) 71 if err != nil { 72 return err 73 } 74 defer conn.Close() 75 gtid, rotatedLog, err := conn.streamWithSnapshot(rs.ctx, rs.tableName.String(), rs.query) 76 if rotatedLog { 77 rs.vse.vstreamerFlushedBinlogs.Add(1) 78 } 79 if err != nil { 80 return err 81 } 82 83 // first call the callback with the fields 84 flds, err := conn.Fields() 85 if err != nil { 86 return err 87 } 88 89 err = rs.send(&binlogdatapb.VStreamResultsResponse{ 90 Fields: flds, 91 Gtid: gtid, 92 }) 93 if err != nil { 94 return fmt.Errorf("stream send error: %v", err) 95 } 96 97 response := &binlogdatapb.VStreamResultsResponse{} 98 byteCount := 0 99 for { 100 select { 101 case <-rs.ctx.Done(): 102 return fmt.Errorf("stream ended: %v", rs.ctx.Err()) 103 default: 104 } 105 106 // check throttler. 107 if !rs.vse.throttlerClient.ThrottleCheckOKOrWait(rs.ctx) { 108 continue 109 } 110 111 row, err := conn.FetchNext(nil) 112 if err != nil { 113 return err 114 } 115 if row == nil { 116 break 117 } 118 response.Rows = append(response.Rows, sqltypes.RowToProto3(row)) 119 for _, s := range row { 120 byteCount += s.Len() 121 } 122 123 if rs.pktsize.ShouldSend(byteCount) { 124 rs.vse.resultStreamerNumRows.Add(int64(len(response.Rows))) 125 rs.vse.resultStreamerNumPackets.Add(int64(1)) 126 startSend := time.Now() 127 err = rs.send(response) 128 if err != nil { 129 return err 130 } 131 rs.pktsize.Record(byteCount, time.Since(startSend)) 132 // empty the rows so we start over, but we keep the 133 // same capacity 134 response.Rows = response.Rows[:0] 135 byteCount = 0 136 } 137 } 138 139 if len(response.Rows) > 0 { 140 rs.vse.resultStreamerNumRows.Add(int64(len(response.Rows))) 141 err = rs.send(response) 142 if err != nil { 143 return err 144 } 145 } 146 147 return nil 148 }