vitess.io/vitess@v0.16.2/go/vt/vtctld/tablet_data_test.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 vtctld
    18  
    19  import (
    20  	"io"
    21  	"sync"
    22  	"testing"
    23  	"time"
    24  
    25  	"context"
    26  
    27  	"google.golang.org/protobuf/proto"
    28  
    29  	"vitess.io/vitess/go/vt/logutil"
    30  	querypb "vitess.io/vitess/go/vt/proto/query"
    31  	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
    32  	"vitess.io/vitess/go/vt/topo/memorytopo"
    33  	"vitess.io/vitess/go/vt/vttablet/grpcqueryservice"
    34  	"vitess.io/vitess/go/vt/vttablet/queryservice"
    35  	"vitess.io/vitess/go/vt/vttablet/queryservice/fakes"
    36  	"vitess.io/vitess/go/vt/vttablet/tmclient"
    37  	"vitess.io/vitess/go/vt/wrangler"
    38  	"vitess.io/vitess/go/vt/wrangler/testlib"
    39  )
    40  
    41  // streamHealthTabletServer is a local QueryService implementation to support the tests
    42  type streamHealthTabletServer struct {
    43  	queryservice.QueryService
    44  	t *testing.T
    45  
    46  	// streamHealthMutex protects all the following fields
    47  	streamHealthMutex sync.Mutex
    48  	streamHealthIndex int
    49  	streamHealthMap   map[int]chan<- *querypb.StreamHealthResponse
    50  }
    51  
    52  func newStreamHealthTabletServer(t *testing.T) *streamHealthTabletServer {
    53  	return &streamHealthTabletServer{
    54  		QueryService:    fakes.ErrorQueryService,
    55  		t:               t,
    56  		streamHealthMap: make(map[int]chan<- *querypb.StreamHealthResponse),
    57  	}
    58  }
    59  
    60  func (s *streamHealthTabletServer) StreamHealth(ctx context.Context, callback func(*querypb.StreamHealthResponse) error) error {
    61  	id, ch := s.streamHealthRegister()
    62  	defer s.streamHealthUnregister(id)
    63  	for shr := range ch {
    64  		if err := callback(shr); err != nil {
    65  			if err == io.EOF {
    66  				return nil
    67  			}
    68  			return err
    69  		}
    70  	}
    71  	return nil
    72  }
    73  
    74  func (s *streamHealthTabletServer) streamHealthRegister() (id int, ch chan *querypb.StreamHealthResponse) {
    75  	s.streamHealthMutex.Lock()
    76  	defer s.streamHealthMutex.Unlock()
    77  
    78  	id = s.streamHealthIndex
    79  	s.streamHealthIndex++
    80  	ch = make(chan *querypb.StreamHealthResponse, 10)
    81  	s.streamHealthMap[id] = ch
    82  	return id, ch
    83  }
    84  
    85  func (s *streamHealthTabletServer) streamHealthUnregister(id int) error {
    86  	s.streamHealthMutex.Lock()
    87  	defer s.streamHealthMutex.Unlock()
    88  
    89  	delete(s.streamHealthMap, id)
    90  	return nil
    91  }
    92  
    93  // BroadcastHealth will broadcast the current health to all listeners
    94  func (s *streamHealthTabletServer) BroadcastHealth() {
    95  	shr := &querypb.StreamHealthResponse{
    96  		TabletExternallyReparentedTimestamp: 42,
    97  		RealtimeStats: &querypb.RealtimeStats{
    98  			HealthError:           "testHealthError",
    99  			ReplicationLagSeconds: 72,
   100  			CpuUsage:              1.1,
   101  		},
   102  	}
   103  
   104  	s.streamHealthMutex.Lock()
   105  	defer s.streamHealthMutex.Unlock()
   106  	for _, c := range s.streamHealthMap {
   107  		c <- shr
   108  	}
   109  }
   110  
   111  func TestTabletData(t *testing.T) {
   112  	ts := memorytopo.NewServer("cell1", "cell2")
   113  	wr := wrangler.New(logutil.NewConsoleLogger(), ts, tmclient.NewTabletManagerClient())
   114  
   115  	if err := ts.CreateKeyspace(context.Background(), "ks", &topodatapb.Keyspace{}); err != nil {
   116  		t.Fatalf("CreateKeyspace failed: %v", err)
   117  	}
   118  
   119  	tablet1 := testlib.NewFakeTablet(t, wr, "cell1", 0, topodatapb.TabletType_PRIMARY, nil, testlib.TabletKeyspaceShard(t, "ks", "-80"))
   120  	shsq := newStreamHealthTabletServer(t)
   121  	grpcqueryservice.Register(tablet1.RPCServer, shsq)
   122  	tablet1.StartActionLoop(t, wr)
   123  	defer tablet1.StopActionLoop(t)
   124  
   125  	thc := newTabletHealthCache(ts)
   126  
   127  	// Keep broadcasting until the first result goes through.
   128  	stop := make(chan struct{})
   129  	go func() {
   130  		for {
   131  			select {
   132  			case <-stop:
   133  				return
   134  			default:
   135  				shsq.BroadcastHealth()
   136  			}
   137  		}
   138  	}()
   139  
   140  	// Start streaming and wait for the first result.
   141  	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
   142  	result, err := thc.Get(ctx, tablet1.Tablet.Alias)
   143  	cancel()
   144  	close(stop)
   145  
   146  	if err != nil {
   147  		t.Fatalf("thc.Get failed: %v", err)
   148  	}
   149  
   150  	stats := &querypb.RealtimeStats{
   151  		HealthError:           "testHealthError",
   152  		ReplicationLagSeconds: 72,
   153  		CpuUsage:              1.1,
   154  	}
   155  	if got, want := result.RealtimeStats, stats; !proto.Equal(got, want) {
   156  		t.Errorf("RealtimeStats = %#v, want %#v", got, want)
   157  	}
   158  }