vitess.io/vitess@v0.16.2/go/vt/vttablet/queryservice/fakes/stream_health_query_service.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 fakes
    18  
    19  import (
    20  	"context"
    21  
    22  	"google.golang.org/protobuf/proto"
    23  
    24  	"vitess.io/vitess/go/sqltypes"
    25  	querypb "vitess.io/vitess/go/vt/proto/query"
    26  	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
    27  	"vitess.io/vitess/go/vt/vttablet/queryservice"
    28  )
    29  
    30  const (
    31  	// DefaultReplicationLagSeconds is the default MySQL replication lag which is
    32  	// reported in all faked stream health responses.
    33  	DefaultReplicationLagSeconds uint32 = 1
    34  )
    35  
    36  // StreamHealthQueryService is a QueryService implementation which allows to
    37  // send custom StreamHealthResponse messages by adding them to a channel.
    38  // Note that it only works with one connected client because messages going
    39  // into "healthResponses" are not duplicated to all clients.
    40  //
    41  // If you want to override other QueryService methods, embed this struct
    42  // as anonymous field in your own QueryService fake.
    43  type StreamHealthQueryService struct {
    44  	queryservice.QueryService
    45  	healthResponses chan *querypb.StreamHealthResponse
    46  	target          *querypb.Target
    47  }
    48  
    49  var _ queryservice.QueryService = (*StreamHealthQueryService)(nil)
    50  
    51  // NewStreamHealthQueryService creates a new fake query service for the target.
    52  func NewStreamHealthQueryService(target *querypb.Target) *StreamHealthQueryService {
    53  	return &StreamHealthQueryService{
    54  		QueryService:    ErrorQueryService,
    55  		healthResponses: make(chan *querypb.StreamHealthResponse, 1000),
    56  		target:          proto.Clone(target).(*querypb.Target),
    57  	}
    58  }
    59  
    60  // Begin implemented as a no op
    61  func (q *StreamHealthQueryService) Begin(ctx context.Context, target *querypb.Target, options *querypb.ExecuteOptions) (queryservice.TransactionState, error) {
    62  	return queryservice.TransactionState{}, nil
    63  }
    64  
    65  // Execute implemented as a no op
    66  func (q *StreamHealthQueryService) Execute(ctx context.Context, target *querypb.Target, sql string, bindVariables map[string]*querypb.BindVariable, transactionID, reservedID int64, options *querypb.ExecuteOptions) (*sqltypes.Result, error) {
    67  	return &sqltypes.Result{}, nil
    68  }
    69  
    70  // StreamHealth implements the QueryService interface.
    71  // It sends all queued and future healthResponses to the connected client e.g.
    72  // the healthcheck module.
    73  func (q *StreamHealthQueryService) StreamHealth(ctx context.Context, callback func(*querypb.StreamHealthResponse) error) error {
    74  	for shr := range q.healthResponses {
    75  		callback(shr) // nolint:errcheck
    76  	}
    77  	return nil
    78  }
    79  
    80  // AddDefaultHealthResponse adds a faked health response to the buffer channel.
    81  // The response will have default values typical for a healthy tablet.
    82  func (q *StreamHealthQueryService) AddDefaultHealthResponse() {
    83  	q.healthResponses <- &querypb.StreamHealthResponse{
    84  		Target:  proto.Clone(q.target).(*querypb.Target),
    85  		Serving: true,
    86  		RealtimeStats: &querypb.RealtimeStats{
    87  			ReplicationLagSeconds: DefaultReplicationLagSeconds,
    88  		},
    89  	}
    90  }
    91  
    92  // AddHealthResponseWithQPS adds a faked health response to the buffer channel.
    93  // Only "qps" is different in this message.
    94  func (q *StreamHealthQueryService) AddHealthResponseWithQPS(qps float64) {
    95  	q.healthResponses <- &querypb.StreamHealthResponse{
    96  		Target:  proto.Clone(q.target).(*querypb.Target),
    97  		Serving: true,
    98  		RealtimeStats: &querypb.RealtimeStats{
    99  			Qps:                   qps,
   100  			ReplicationLagSeconds: DefaultReplicationLagSeconds,
   101  		},
   102  	}
   103  }
   104  
   105  // AddHealthResponseWithReplicationLag adds a faked health response to the
   106  // buffer channel. Only "replication_lag_seconds" is different in this message.
   107  func (q *StreamHealthQueryService) AddHealthResponseWithReplicationLag(replicationLag uint32) {
   108  	q.healthResponses <- &querypb.StreamHealthResponse{
   109  		Target:  proto.Clone(q.target).(*querypb.Target),
   110  		Serving: true,
   111  		RealtimeStats: &querypb.RealtimeStats{
   112  			ReplicationLagSeconds: replicationLag,
   113  		},
   114  	}
   115  }
   116  
   117  // AddHealthResponseWithNotServing adds a faked health response to the
   118  // buffer channel. Only "Serving" is different in this message.
   119  func (q *StreamHealthQueryService) AddHealthResponseWithNotServing() {
   120  	q.healthResponses <- &querypb.StreamHealthResponse{
   121  		Target:  proto.Clone(q.target).(*querypb.Target),
   122  		Serving: false,
   123  		RealtimeStats: &querypb.RealtimeStats{
   124  			ReplicationLagSeconds: DefaultReplicationLagSeconds,
   125  		},
   126  	}
   127  }
   128  
   129  // UpdateType changes the type of the query service.
   130  // Only newly sent health messages will use the new type.
   131  func (q *StreamHealthQueryService) UpdateType(tabletType topodatapb.TabletType) {
   132  	q.target.TabletType = tabletType
   133  }