vitess.io/vitess@v0.16.2/go/vt/vttablet/endtoend/framework/eventcatcher.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 framework
    18  
    19  import (
    20  	"errors"
    21  	"time"
    22  
    23  	"vitess.io/vitess/go/streamlog"
    24  	"vitess.io/vitess/go/vt/vttablet/tabletserver"
    25  	"vitess.io/vitess/go/vt/vttablet/tabletserver/tabletenv"
    26  )
    27  
    28  // TxCatcher allows you to capture and fetch transactions that are being
    29  // executed by TabletServer.
    30  type TxCatcher struct {
    31  	catcher *eventCatcher
    32  }
    33  
    34  // NewTxCatcher sets up the capture and returns a new TxCatcher.
    35  // You must call Close when done.
    36  func NewTxCatcher() TxCatcher {
    37  	return TxCatcher{catcher: newEventCatcher(tabletenv.TxLogger)}
    38  }
    39  
    40  // Close closes the TxCatcher.
    41  func (tc *TxCatcher) Close() {
    42  	tc.catcher.Close()
    43  }
    44  
    45  // Next fetches the next captured transaction.
    46  // If the wait is longer than one second, it returns an error.
    47  func (tc *TxCatcher) Next() (*tabletserver.StatefulConnection, error) {
    48  	event, err := tc.catcher.next()
    49  	if err != nil {
    50  		return nil, err
    51  	}
    52  	return event.(*tabletserver.StatefulConnection), nil
    53  }
    54  
    55  // QueryCatcher allows you to capture and fetch queries that are being
    56  // executed by TabletServer.
    57  type QueryCatcher struct {
    58  	catcher *eventCatcher
    59  }
    60  
    61  // NewQueryCatcher sets up the capture and returns a QueryCatcher.
    62  // You must call Close when done.
    63  func NewQueryCatcher() QueryCatcher {
    64  	return QueryCatcher{catcher: newEventCatcher(tabletenv.StatsLogger)}
    65  }
    66  
    67  // Close closes the QueryCatcher.
    68  func (qc *QueryCatcher) Close() {
    69  	qc.catcher.Close()
    70  }
    71  
    72  // Next fetches the next captured query.
    73  // If the wait is longer than one second, it returns an error.
    74  func (qc *QueryCatcher) Next() (*tabletenv.LogStats, error) {
    75  	event, err := qc.catcher.next()
    76  	if err != nil {
    77  		return nil, err
    78  	}
    79  	return event.(*tabletenv.LogStats), nil
    80  }
    81  
    82  type eventCatcher struct {
    83  	start   time.Time
    84  	logger  *streamlog.StreamLogger
    85  	in, out chan any
    86  }
    87  
    88  func newEventCatcher(logger *streamlog.StreamLogger) *eventCatcher {
    89  	catcher := &eventCatcher{
    90  		start:  time.Now(),
    91  		logger: logger,
    92  		in:     logger.Subscribe("endtoend"),
    93  		out:    make(chan any, 20),
    94  	}
    95  	go func() {
    96  		for event := range catcher.in {
    97  			endTime := event.(interface {
    98  				EventTime() time.Time
    99  			}).EventTime()
   100  			if endTime.Before(catcher.start) {
   101  				continue
   102  			}
   103  			catcher.out <- event
   104  		}
   105  		close(catcher.out)
   106  	}()
   107  	return catcher
   108  }
   109  
   110  // Close closes the eventCatcher.
   111  func (catcher *eventCatcher) Close() {
   112  	catcher.logger.Unsubscribe(catcher.in)
   113  	close(catcher.in)
   114  }
   115  
   116  func (catcher *eventCatcher) next() (any, error) {
   117  	tmr := time.NewTimer(5 * time.Second)
   118  	defer tmr.Stop()
   119  	for {
   120  		select {
   121  		case event := <-catcher.out:
   122  			return event, nil
   123  		case <-tmr.C:
   124  			return nil, errors.New("error waiting for query event")
   125  		}
   126  	}
   127  }