github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/sql/physicalplan/replicaoracle/oracle_test.go (about)

     1  // Copyright 2019 The Cockroach Authors.
     2  //
     3  // Use of this software is governed by the Business Source License
     4  // included in the file licenses/BSL.txt.
     5  //
     6  // As of the Change Date specified in that file, in accordance with
     7  // the Business Source License, use of this software will be governed
     8  // by the Apache License, Version 2.0, included in the file
     9  // licenses/APL.txt.
    10  
    11  package replicaoracle
    12  
    13  import (
    14  	"context"
    15  	"fmt"
    16  	"strings"
    17  	"testing"
    18  	"time"
    19  
    20  	"github.com/cockroachdb/cockroach/pkg/config/zonepb"
    21  	"github.com/cockroachdb/cockroach/pkg/gossip"
    22  	"github.com/cockroachdb/cockroach/pkg/roachpb"
    23  	"github.com/cockroachdb/cockroach/pkg/rpc"
    24  	"github.com/cockroachdb/cockroach/pkg/util"
    25  	"github.com/cockroachdb/cockroach/pkg/util/hlc"
    26  	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
    27  	"github.com/cockroachdb/cockroach/pkg/util/metric"
    28  	"github.com/cockroachdb/cockroach/pkg/util/stop"
    29  )
    30  
    31  // TestRandomOracle defeats TestUnused for RandomChoice.
    32  func TestRandomOracle(t *testing.T) {
    33  	_ = NewOracleFactory(RandomChoice, Config{})
    34  }
    35  
    36  func TestClosest(t *testing.T) {
    37  	defer leaktest.AfterTest(t)()
    38  	stopper := stop.NewStopper()
    39  	defer stopper.Stop(context.Background())
    40  	g, _ := makeGossip(t, stopper)
    41  	nd, _ := g.GetNodeDescriptor(1)
    42  	of := NewOracleFactory(ClosestChoice, Config{
    43  		Gossip:   g,
    44  		NodeDesc: *nd,
    45  	})
    46  	of.(*closestOracle).latencyFunc = func(s string) (time.Duration, bool) {
    47  		if strings.HasSuffix(s, "2") {
    48  			return time.Nanosecond, true
    49  		}
    50  		return time.Millisecond, true
    51  	}
    52  	o := of.Oracle(nil)
    53  	info, err := o.ChoosePreferredReplica(context.Background(), roachpb.RangeDescriptor{
    54  		InternalReplicas: []roachpb.ReplicaDescriptor{
    55  			{NodeID: 1, StoreID: 1},
    56  			{NodeID: 2, StoreID: 2},
    57  			{NodeID: 3, StoreID: 3},
    58  		},
    59  	}, QueryState{})
    60  	if err != nil {
    61  		t.Fatalf("Failed to choose closest replica: %v", err)
    62  	}
    63  	if info.NodeID != 2 {
    64  		t.Fatalf("Failed to choose node 2, got %v", info.NodeID)
    65  	}
    66  }
    67  
    68  func makeGossip(t *testing.T, stopper *stop.Stopper) (*gossip.Gossip, *hlc.Clock) {
    69  	clock := hlc.NewClock(hlc.UnixNano, time.Nanosecond)
    70  	rpcContext := rpc.NewInsecureTestingContext(clock, stopper)
    71  	server := rpc.NewServer(rpcContext)
    72  
    73  	const nodeID = 1
    74  	g := gossip.NewTest(nodeID, rpcContext, server, stopper, metric.NewRegistry(), zonepb.DefaultZoneConfigRef())
    75  	if err := g.SetNodeDescriptor(newNodeDesc(nodeID)); err != nil {
    76  		t.Fatal(err)
    77  	}
    78  	if err := g.AddInfo(gossip.KeySentinel, nil, time.Hour); err != nil {
    79  		t.Fatal(err)
    80  	}
    81  	for i := roachpb.NodeID(2); i <= 3; i++ {
    82  		err := g.AddInfoProto(gossip.MakeNodeIDKey(i), newNodeDesc(i), gossip.NodeDescriptorTTL)
    83  		if err != nil {
    84  			t.Fatal(err)
    85  		}
    86  	}
    87  	return g, clock
    88  }
    89  
    90  func newNodeDesc(nodeID roachpb.NodeID) *roachpb.NodeDescriptor {
    91  	return &roachpb.NodeDescriptor{
    92  		NodeID:  nodeID,
    93  		Address: util.MakeUnresolvedAddr("tcp", fmt.Sprintf("invalid.invalid:%d", nodeID)),
    94  	}
    95  }