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  }