mirror of
https://github.com/featurebasedb/featurebase.git
synced 2026-08-28 10:54:59 +00:00
815 lines
25 KiB
Go
815 lines
25 KiB
Go
// Copyright 2022 Molecula Corp. (DBA FeatureBase).
|
|
// SPDX-License-Identifier: Apache-2.0
|
|
package test
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"math"
|
|
"sort"
|
|
"strings"
|
|
"sync"
|
|
"testing"
|
|
"time"
|
|
"unicode"
|
|
|
|
pilosa "github.com/featurebasedb/featurebase/v3"
|
|
"github.com/featurebasedb/featurebase/v3/api/client"
|
|
"github.com/featurebasedb/featurebase/v3/disco"
|
|
"github.com/featurebasedb/featurebase/v3/etcd"
|
|
"github.com/featurebasedb/featurebase/v3/logger"
|
|
"github.com/featurebasedb/featurebase/v3/proto"
|
|
"github.com/featurebasedb/featurebase/v3/server"
|
|
"github.com/featurebasedb/featurebase/v3/storage"
|
|
"github.com/pkg/errors"
|
|
"golang.org/x/sync/errgroup"
|
|
)
|
|
|
|
// Cluster represents a per-test wrapper of a "real" cluster.
|
|
// Individual tests which request a cluster can get one of these
|
|
// back, wrapped with their test name.
|
|
type Cluster struct {
|
|
*ShareableCluster
|
|
tb testing.TB
|
|
indexName string
|
|
indexBytes []byte
|
|
}
|
|
|
|
// Idx produces an index name suitable for this test's
|
|
// cluster, using the name itself, or the concatenation
|
|
// of that name and any provided strings, joined by underscores.
|
|
func (c *Cluster) Idx(of ...string) string {
|
|
if len(of) == 0 {
|
|
return c.indexName
|
|
}
|
|
for i := range of {
|
|
of[i] = indexName(of[i])
|
|
}
|
|
return c.indexName + strings.Join(of, "_")
|
|
}
|
|
|
|
// Format implements Formatter, allowing us to use clusters in
|
|
// format strings to handle the extremely common problem of "I
|
|
// want to embed the index name in this here string". Because we
|
|
// are horrible criminals, we append the rune, so you can use
|
|
// %i and %j to get index-name-plus-i and index-name-plus-j,
|
|
// respectively. We do assume that the rune works as a plain byte,
|
|
// though.
|
|
func (c *Cluster) Format(state fmt.State, r rune) {
|
|
c.indexBytes = append(c.indexBytes[:0], c.indexName...)
|
|
c.indexBytes = append(c.indexBytes, byte(r))
|
|
_, _ = state.Write(c.indexBytes)
|
|
}
|
|
|
|
func (c *Cluster) Close() error {
|
|
if c.shared {
|
|
// We're not going to actually close the cluster, BUT, we do want to
|
|
// delete and close all the indexes we might have just made.
|
|
h := c.GetHolder(0)
|
|
indexes := h.Indexes()
|
|
var lastErr error
|
|
api := c.Nodes[0].API
|
|
for _, idx := range indexes {
|
|
name := idx.Name()
|
|
if strings.HasPrefix(name, c.indexName) {
|
|
err := api.DeleteIndex(context.Background(), name)
|
|
if err != nil {
|
|
lastErr = err
|
|
}
|
|
}
|
|
}
|
|
return lastErr
|
|
}
|
|
if !c.started {
|
|
return errors.New("cluster not started yet")
|
|
}
|
|
return c.ShareableCluster.Close()
|
|
}
|
|
|
|
func (c *Cluster) Start() error {
|
|
if c.started {
|
|
return errors.New("cluster already started")
|
|
}
|
|
c.started = true
|
|
return c.ShareableCluster.Start()
|
|
}
|
|
|
|
// ShareableCluster represents a featurebase cluster (multiple Command instances)
|
|
// without test-specific overhead.
|
|
type ShareableCluster struct {
|
|
Nodes []*Command
|
|
started bool
|
|
shared bool
|
|
}
|
|
|
|
// Query executes an API.Query through one of the cluster's node's API. It fails
|
|
// the test if there is an error.
|
|
func (c *ShareableCluster) Query(t testing.TB, index, query string) pilosa.QueryResponse {
|
|
t.Helper()
|
|
if len(c.Nodes) == 0 {
|
|
t.Fatal("must have at least one node in cluster to query")
|
|
}
|
|
|
|
return c.GetPrimary().QueryAPI(t, &pilosa.QueryRequest{Index: index, Query: query})
|
|
}
|
|
|
|
// QueryHTTP executes a PQL query through the HTTP endpoint. It fails
|
|
// the test for explicit errors, but returns an error which has the
|
|
// response body if the HTTP call returns a non-OK status.
|
|
func (c *ShareableCluster) QueryHTTP(t testing.TB, index, query string) (string, error) {
|
|
t.Helper()
|
|
if len(c.Nodes) == 0 {
|
|
t.Fatal("must have at least one node in cluster to query")
|
|
}
|
|
|
|
return c.GetPrimary().Query(t, index, "", query)
|
|
}
|
|
|
|
// QueryGRPC executes a PQL query through the GRPC endpoint. It fails the
|
|
// test if there is an error.
|
|
func (c *ShareableCluster) QueryGRPC(t testing.TB, index, query string) *proto.TableResponse {
|
|
t.Helper()
|
|
if len(c.Nodes) == 0 {
|
|
t.Fatal("must have at least one node in cluster to query")
|
|
}
|
|
|
|
grpcClient, err := client.NewGRPCClient([]string{fmt.Sprintf("%s:%d", c.GetPrimary().Server.GRPCURI().Host, c.GetPrimary().Server.GRPCURI().Port)}, nil, logger.NopLogger)
|
|
if err != nil {
|
|
t.Fatalf("getting GRPC client: %v", err)
|
|
}
|
|
defer grpcClient.Close()
|
|
|
|
tableResp, err := grpcClient.QueryUnary(context.Background(), index, query)
|
|
if err != nil {
|
|
t.Fatalf("querying unary: %v", err)
|
|
}
|
|
|
|
return tableResp
|
|
}
|
|
|
|
// GetIdleNode gets the node at the given index. This method is used (instead of
|
|
// `GetNode()`) when the cluster has yet to be started. In that case, etcd has
|
|
// not assigned each node an ID, and therefore the nodes are not in their final,
|
|
// sorted order. In other words, this method can only be used to retrieve a node
|
|
// when order doesn't matter. An example is if you need to do something like
|
|
// this:
|
|
//
|
|
// c.GetNode(0).Config.Cluster.ReplicaN = 2
|
|
// c.GetNode(1).Config.Cluster.ReplicaN = 2
|
|
//
|
|
// In this example, the test needs the replication factor to be set to 2 before
|
|
// starting; it's ok to reference each node by its index in the pre-sorted node
|
|
// list. It's also safe to use this method after `MustRunCluster()` if the
|
|
// cluster contains only one node.
|
|
func (c *ShareableCluster) GetIdleNode(n int) *Command {
|
|
return c.Nodes[n]
|
|
}
|
|
|
|
// GetNode gets the node at the given index; this method assumes the cluster has
|
|
// already been started. Because the node IDs are assigned randomly, they can be
|
|
// in an order that does not align with the test's expectations. For example, a
|
|
// test might create a 3-node cluster and retrieve them using `GetNode(0)`,
|
|
// `GetNode(1)`, and `GetNode(2)` respectively. But if the node IDs are `456`,
|
|
// `123`, `789`, then we actually want `GetNode(0)` to return `c.Nodes[1]`, and
|
|
// `GetNode(1)` to return `c.Nodes[0]`. This method looks at all the node IDs,
|
|
// sorts them, and then returns the node that the test expects.
|
|
func (c *ShareableCluster) GetNode(n int) *Command {
|
|
// Put all the node IDs into a list to be sorted.
|
|
ids := make([]nodePlace, len(c.Nodes))
|
|
for i := range c.Nodes {
|
|
ids[i].id = c.Nodes[i].ID()
|
|
ids[i].idx = i
|
|
}
|
|
|
|
// Sort the list.
|
|
sort.SliceStable(ids, func(i, j int) bool {
|
|
return ids[i].id < ids[j].id
|
|
})
|
|
|
|
// Return the node which is at the given position in the sorted list.
|
|
return c.Nodes[ids[n].idx]
|
|
}
|
|
|
|
// GetPrimary gets the node which has been determined to be the primary.
|
|
// This used to be node0 in tests, but since implementing etcd, the primary
|
|
// can be any node in the cluster, so we have to use this method in tests which
|
|
// need to act on the primary.
|
|
func (c *ShareableCluster) GetPrimary() *Command {
|
|
for _, n := range c.Nodes {
|
|
if n.IsPrimary() {
|
|
return n
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// GetNonPrimary gets first first non-primary node in the list of nodes.
|
|
func (c *ShareableCluster) GetNonPrimary() *Command {
|
|
for _, n := range c.Nodes {
|
|
if !n.IsPrimary() {
|
|
return n
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// GetNonPrimaries gets all nodes except the primary.
|
|
func (c *ShareableCluster) GetNonPrimaries() []*Command {
|
|
rtn := make([]*Command, 0)
|
|
for _, n := range c.Nodes {
|
|
if !n.IsPrimary() {
|
|
rtn = append(rtn, n)
|
|
}
|
|
}
|
|
return rtn
|
|
}
|
|
|
|
// nodePlace represents a node's ID and its index into the c.Nodes slice.
|
|
type nodePlace struct {
|
|
id string
|
|
idx int
|
|
}
|
|
|
|
func (c *ShareableCluster) GetHolder(n int) *Holder {
|
|
return &Holder{Holder: c.GetNode(n).Server.Holder()}
|
|
}
|
|
|
|
func (c *ShareableCluster) Len() int {
|
|
return len(c.Nodes)
|
|
}
|
|
|
|
func (c *ShareableCluster) ImportBitsWithTimestamp(t testing.TB, index, field string, rowcols [][2]uint64, timestamps []int64) {
|
|
t.Helper()
|
|
byShard := make(map[uint64][][2]uint64)
|
|
byShardTs := make(map[uint64][]int64)
|
|
for i, rowcol := range rowcols {
|
|
shard := rowcol[1] / pilosa.ShardWidth
|
|
byShard[shard] = append(byShard[shard], rowcol)
|
|
if len(timestamps) > 0 {
|
|
byShardTs[shard] = append(byShardTs[shard], timestamps[i])
|
|
}
|
|
}
|
|
|
|
for shard, bits := range byShard {
|
|
rowIDs := make([]uint64, len(bits))
|
|
colIDs := make([]uint64, len(bits))
|
|
for i, bit := range bits {
|
|
rowIDs[i] = bit[0]
|
|
colIDs[i] = bit[1]
|
|
}
|
|
nodes, err := c.GetPrimary().API.ShardNodes(context.Background(), index, shard)
|
|
if err != nil {
|
|
t.Fatalf("getting shard nodes: %v", err)
|
|
}
|
|
// TODO won't be necessary to do all nodes once that works hits
|
|
// (travis) this TODO is not clear to me, but I think it's
|
|
// suggesting that elsewhere we would support importing to a
|
|
// single node, regardless of where the data ends up.
|
|
for _, node := range nodes {
|
|
for _, com := range c.Nodes {
|
|
if com.API.Node().ID != node.ID {
|
|
continue
|
|
}
|
|
func() {
|
|
qcx := com.API.Txf().NewQcx()
|
|
defer qcx.Abort()
|
|
if len(timestamps) == 0 {
|
|
err := com.API.Import(context.Background(), qcx, &pilosa.ImportRequest{
|
|
Index: index,
|
|
Field: field,
|
|
Shard: shard,
|
|
RowIDs: rowIDs,
|
|
ColumnIDs: colIDs,
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("importing data: %v", err)
|
|
}
|
|
} else {
|
|
ts := byShardTs[shard]
|
|
err := com.API.Import(context.Background(), qcx, &pilosa.ImportRequest{
|
|
Index: index,
|
|
Field: field,
|
|
Shard: shard,
|
|
RowIDs: rowIDs,
|
|
ColumnIDs: colIDs,
|
|
Timestamps: ts,
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("importing data: %v", err)
|
|
}
|
|
|
|
}
|
|
}()
|
|
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
func (c *ShareableCluster) ImportBits(t testing.TB, index, field string, rowcols [][2]uint64) {
|
|
var noTime []int64
|
|
c.ImportBitsWithTimestamp(t, index, field, rowcols, noTime)
|
|
}
|
|
|
|
// ImportKeyKey imports data into an index where both the index and
|
|
// the field are using string keys.
|
|
func (c *ShareableCluster) ImportKeyKey(t testing.TB, index, field string, valAndRecKeys [][2]string) {
|
|
t.Helper()
|
|
importRequest := &pilosa.ImportRequest{
|
|
Index: index,
|
|
Field: field,
|
|
RowKeys: make([]string, len(valAndRecKeys)),
|
|
ColumnKeys: make([]string, len(valAndRecKeys)),
|
|
}
|
|
for i, vk := range valAndRecKeys {
|
|
importRequest.RowKeys[i] = vk[0]
|
|
importRequest.ColumnKeys[i] = vk[1]
|
|
}
|
|
qcx := c.GetPrimary().API.Txf().NewQcx()
|
|
defer qcx.Abort()
|
|
err := c.GetPrimary().API.Import(context.Background(), qcx, importRequest)
|
|
if err != nil {
|
|
t.Fatalf("importing keykey data: %v", err)
|
|
}
|
|
}
|
|
|
|
// TimeQuantumKey is a string key and a string+key value
|
|
type TimeQuantumKey struct {
|
|
RowKey string
|
|
ColKey string
|
|
Ts int64
|
|
}
|
|
|
|
// ImportTimeQuantumKey imports data into an index where the index is keyd
|
|
// and the field is a time-quantum
|
|
func (c *ShareableCluster) ImportTimeQuantumKey(t testing.TB, index, field string, entries []TimeQuantumKey) {
|
|
t.Helper()
|
|
importRequest := &pilosa.ImportRequest{
|
|
Index: index,
|
|
Field: field,
|
|
RowKeys: make([]string, len(entries)),
|
|
ColumnKeys: make([]string, len(entries)),
|
|
Timestamps: make([]int64, len(entries)),
|
|
}
|
|
for i, entry := range entries {
|
|
importRequest.ColumnKeys[i] = entry.ColKey
|
|
importRequest.RowKeys[i] = entry.RowKey
|
|
importRequest.Timestamps[i] = entry.Ts
|
|
|
|
}
|
|
qcx := c.GetPrimary().API.Txf().NewQcx()
|
|
defer qcx.Abort()
|
|
err := c.GetPrimary().API.Import(context.Background(), qcx, importRequest)
|
|
if err != nil {
|
|
t.Fatalf("importing keykey data: %v", err)
|
|
}
|
|
}
|
|
|
|
// IntKey is a string key and a signed integer value.
|
|
type IntKey struct {
|
|
Val int64
|
|
Key string
|
|
}
|
|
|
|
// ImportIntKey imports int data into an index which uses string keys.
|
|
func (c *ShareableCluster) ImportIntKey(t testing.TB, index, field string, pairs []IntKey) {
|
|
t.Helper()
|
|
importRequest := &pilosa.ImportValueRequest{
|
|
Index: index,
|
|
Field: field,
|
|
Shard: math.MaxUint64,
|
|
ColumnKeys: make([]string, len(pairs)),
|
|
Values: make([]int64, len(pairs)),
|
|
}
|
|
for i, pair := range pairs {
|
|
importRequest.Values[i] = pair.Val
|
|
importRequest.ColumnKeys[i] = pair.Key
|
|
}
|
|
qcx := c.GetPrimary().API.Txf().NewQcx()
|
|
defer qcx.Abort()
|
|
if err := c.GetPrimary().API.ImportValue(context.Background(), qcx, importRequest); err != nil {
|
|
t.Fatalf("importing IntKey data: %v", err)
|
|
}
|
|
}
|
|
|
|
type IntID struct {
|
|
Val int64
|
|
ID uint64
|
|
}
|
|
|
|
// ImportIntID imports data into an int field in an unkeyed index.
|
|
func (c *ShareableCluster) ImportIntID(t testing.TB, index, field string, pairs []IntID) {
|
|
t.Helper()
|
|
importRequest := &pilosa.ImportValueRequest{
|
|
Index: index,
|
|
Field: field,
|
|
Shard: math.MaxUint64,
|
|
ColumnIDs: make([]uint64, len(pairs)),
|
|
Values: make([]int64, len(pairs)),
|
|
}
|
|
for i, pair := range pairs {
|
|
importRequest.Values[i] = pair.Val
|
|
importRequest.ColumnIDs[i] = pair.ID
|
|
}
|
|
qcx := c.GetPrimary().API.Txf().NewQcx()
|
|
defer qcx.Abort()
|
|
if err := c.GetPrimary().API.ImportValue(context.Background(), qcx, importRequest); err != nil {
|
|
t.Fatalf("importing IntID data: %v", err)
|
|
}
|
|
}
|
|
|
|
// KeyID represents a key and an ID for importing data into an index
|
|
// and field where one uses string keys and the other does not.
|
|
type KeyID struct {
|
|
Key string
|
|
ID uint64
|
|
}
|
|
|
|
// ImportIDKey imports data into an unkeyed set field in a keyed index.
|
|
func (c *ShareableCluster) ImportIDKey(t testing.TB, index, field string, pairs []KeyID) {
|
|
t.Helper()
|
|
importRequest := &pilosa.ImportRequest{
|
|
Index: index,
|
|
Field: field,
|
|
RowIDs: make([]uint64, len(pairs)),
|
|
ColumnKeys: make([]string, len(pairs)),
|
|
}
|
|
for i, pair := range pairs {
|
|
importRequest.RowIDs[i] = pair.ID
|
|
importRequest.ColumnKeys[i] = pair.Key
|
|
}
|
|
qcx := c.GetPrimary().API.Txf().NewQcx()
|
|
defer qcx.Abort()
|
|
err := c.GetPrimary().API.Import(context.Background(), qcx, importRequest)
|
|
if err != nil {
|
|
t.Fatalf("importing IDKey data: %v", err)
|
|
}
|
|
}
|
|
|
|
// CreateField creates the index (if necessary) and field specified.
|
|
func (c *ShareableCluster) CreateField(t testing.TB, index string, iopts pilosa.IndexOptions, field string, fopts ...pilosa.FieldOption) *pilosa.Field {
|
|
t.Helper()
|
|
idx, err := c.GetPrimary().API.CreateIndex(context.Background(), index, iopts)
|
|
if err != nil && !strings.Contains(err.Error(), "index already exists") {
|
|
t.Fatalf("creating index: %v", err)
|
|
} else if err != nil { // index exists
|
|
idx, err = c.GetPrimary().API.Index(context.Background(), index)
|
|
if err != nil {
|
|
t.Fatalf("getting index: %v", err)
|
|
}
|
|
}
|
|
if idx.Options() != iopts {
|
|
t.Logf("existing index options:\n%v\ndon't match given opts:\n%v\n in pilosa/test.Cluster.CreateField", idx.Options(), iopts)
|
|
}
|
|
|
|
f, err := c.GetPrimary().API.CreateField(context.Background(), index, field, fopts...)
|
|
// we'll assume the field doesn't exist because checking if the options
|
|
// match seems painful.
|
|
if err != nil {
|
|
t.Fatalf("creating field: %v", err)
|
|
}
|
|
return f
|
|
}
|
|
|
|
// Start runs a Cluster
|
|
func (c *ShareableCluster) Start() error {
|
|
var eg errgroup.Group
|
|
for _, cc := range c.Nodes {
|
|
cc := cc
|
|
eg.Go(func() error {
|
|
return cc.Start()
|
|
})
|
|
}
|
|
err := eg.Wait()
|
|
if err != nil {
|
|
return errors.Wrap(err, "starting cluster")
|
|
}
|
|
return c.AwaitState(disco.ClusterStateNormal, 30*time.Second)
|
|
}
|
|
|
|
// Close stops a Cluster
|
|
func (c *ShareableCluster) Close() error {
|
|
for i, cc := range c.Nodes {
|
|
if err := cc.Close(); err != nil {
|
|
return errors.Wrapf(err, "stopping server %d", i)
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (c *ShareableCluster) CloseAndRemoveNonPrimary() error {
|
|
if c.shared {
|
|
return errors.New("can't close-and-remove in shared cluster")
|
|
}
|
|
for i, n := range c.Nodes {
|
|
if !n.IsPrimary() {
|
|
return c.CloseAndRemove(i)
|
|
}
|
|
}
|
|
return errors.New("could not find non-primary node")
|
|
}
|
|
|
|
func (c *ShareableCluster) CloseAndRemove(n int) error {
|
|
if c.shared {
|
|
return errors.New("can't close-and-remove in shared cluster")
|
|
}
|
|
if n < 0 || n >= len(c.Nodes) {
|
|
return fmt.Errorf("close/remove from cluster: index %d out of range (len %d)", n, len(c.Nodes))
|
|
}
|
|
err := c.Nodes[n].Close()
|
|
copy(c.Nodes[n:], c.Nodes[n+1:])
|
|
c.Nodes = c.Nodes[:len(c.Nodes)-1]
|
|
return err
|
|
}
|
|
|
|
// AwaitPrimaryState waits for the cluster primary to reach a specified cluster state.
|
|
// When this happens, we know etcd reached a combination of node states that
|
|
// would imply this cluster state, but some nodes may not have caught up yet;
|
|
// we just test that the coordinator thought the cluster was in the given state.
|
|
func (c *ShareableCluster) AwaitPrimaryState(expectedState disco.ClusterState, timeout time.Duration) error {
|
|
if len(c.Nodes) < 1 {
|
|
return errors.New("can't await coordinator state on an empty cluster")
|
|
}
|
|
primary := c.GetPrimary()
|
|
if primary == nil {
|
|
startTime := time.Now()
|
|
var elapsed time.Duration
|
|
for elapsed = 0; elapsed <= timeout; elapsed = time.Since(startTime) {
|
|
time.Sleep(50 * time.Millisecond)
|
|
primary = c.GetPrimary()
|
|
if primary != nil {
|
|
break
|
|
}
|
|
}
|
|
if primary == nil {
|
|
return errors.New("timed out waiting for cluster to have valid topology")
|
|
}
|
|
// we used up some of our timeout waiting for this
|
|
timeout -= elapsed
|
|
}
|
|
onlyCoordinator := &ShareableCluster{Nodes: []*Command{primary}}
|
|
return onlyCoordinator.AwaitState(expectedState, timeout)
|
|
}
|
|
|
|
// ExceptionalState returns an error if any node in the cluster is not
|
|
// in the expected state.
|
|
func (c *ShareableCluster) ExceptionalState(expectedState disco.ClusterState) error {
|
|
for _, node := range c.Nodes {
|
|
state, err := node.API.State()
|
|
if err != nil || state != expectedState {
|
|
return fmt.Errorf("node %q: state %s: err %v", node.ID(), state, err)
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// AwaitState waits for the whole cluster to reach a specified state.
|
|
func (c *ShareableCluster) AwaitState(expectedState disco.ClusterState, timeout time.Duration) (err error) {
|
|
if len(c.Nodes) < 1 {
|
|
return errors.New("can't await state of an empty cluster")
|
|
}
|
|
startTime := time.Now()
|
|
var elapsed time.Duration
|
|
for elapsed = 0; elapsed <= timeout; elapsed = time.Since(startTime) {
|
|
// Counterintuitive: We're returning if the err *is* nil,
|
|
// meaning we've reached the expected state.
|
|
if err = c.ExceptionalState(expectedState); err == nil {
|
|
return err
|
|
}
|
|
time.Sleep(50 * time.Millisecond)
|
|
}
|
|
return fmt.Errorf("waited %v for cluster to reach state %q: %v",
|
|
elapsed, expectedState, err)
|
|
}
|
|
|
|
// MustNewCluster creates a new cluster or returns an existing one. It never shares
|
|
// a cluster with non-empty opts. If opts contains only one
|
|
// slice of command options, those options are used with every node.
|
|
// If it is empty, default options are used. Otherwise, it must contain size
|
|
// slices of command options, which are used with corresponding nodes.
|
|
func MustNewCluster(tb testing.TB, size int, opts ...[]server.CommandOption) *Cluster {
|
|
if size > 1 && !etcd.AllowCluster() {
|
|
tb.Skip("Testing PLG which does not allow clustering")
|
|
}
|
|
tb.Helper()
|
|
|
|
shareable := len(opts) == 0
|
|
// We want tests to default to using the in-memory translate store, so we
|
|
// prepend opts with that functional option. If a different translate store
|
|
// has been specified, it will override this one.
|
|
opts = prependOpts(opts, size)
|
|
|
|
c, err := newCluster(tb, size, shareable, opts...)
|
|
if err != nil {
|
|
tb.Fatalf("new cluster: %v", err)
|
|
}
|
|
return c
|
|
}
|
|
|
|
// MustUnsharedCluster creates a new cluster. If opts contains only one
|
|
// slice of command options, those options are used with every node.
|
|
// If it is empty, default options are used. Otherwise, it must contain size
|
|
// slices of command options, which are used with corresponding nodes. The
|
|
// new cluster is always unshared.
|
|
func MustUnsharedCluster(tb testing.TB, size int, opts ...[]server.CommandOption) *Cluster {
|
|
if size > 1 && !etcd.AllowCluster() {
|
|
tb.Skip("Testing PLG which does not allow clustering")
|
|
}
|
|
tb.Helper()
|
|
// We want tests to default to using the in-memory translate store, so we
|
|
// prepend opts with that functional option. If a different translate store
|
|
// has been specified, it will override this one.
|
|
opts = prependOpts(opts, size)
|
|
|
|
c, err := newCluster(tb, size, false, opts...)
|
|
if err != nil {
|
|
tb.Fatalf("new cluster: %v", err)
|
|
}
|
|
return c
|
|
}
|
|
|
|
// MustRunUnsharedCluster creates a new cluster. If opts contains only one
|
|
// slice of command options, those options are used with every node.
|
|
// If it is empty, default options are used. Otherwise, it must contain size
|
|
// slices of command options, which are used with corresponding nodes. The
|
|
// new cluster is always unshared. The new cluster is started automatically,
|
|
// or the test is failed.
|
|
func MustRunUnsharedCluster(tb testing.TB, size int, opts ...[]server.CommandOption) *Cluster {
|
|
if size > 1 && !etcd.AllowCluster() {
|
|
tb.Skip("Testing PLG which does not allow clustering")
|
|
}
|
|
tb.Helper()
|
|
// We want tests to default to using the in-memory translate store, so we
|
|
// prepend opts with that functional option. If a different translate store
|
|
// has been specified, it will override this one.
|
|
opts = prependOpts(opts, size)
|
|
|
|
c, err := newCluster(tb, size, false, opts...)
|
|
if err != nil {
|
|
tb.Fatalf("new cluster: %v", err)
|
|
}
|
|
err = c.Start()
|
|
if err != nil {
|
|
tb.Fatalf("starting cluster: %v", err)
|
|
}
|
|
return c
|
|
}
|
|
|
|
type clusterCache struct {
|
|
mu sync.Mutex
|
|
clusters map[int]*ShareableCluster
|
|
}
|
|
|
|
// CleanupClusters calls the close functions on any shared clusters that are
|
|
// still open.
|
|
func (c *clusterCache) CleanupClusters() {
|
|
for k, v := range c.clusters {
|
|
_ = v.Close()
|
|
delete(c.clusters, k)
|
|
}
|
|
}
|
|
|
|
func (c *clusterCache) newCluster(tb testing.TB, size int) (*ShareableCluster, error) {
|
|
c.mu.Lock()
|
|
defer c.mu.Unlock()
|
|
if c := c.clusters[size]; c != nil {
|
|
return c, nil
|
|
}
|
|
// Make a whole-test wrapper so that the cluster we create will use the provided tb
|
|
// for nearly everything, but the call to TempDir inside NewCommand will pick up a
|
|
// persistent directory which outlives the provided TB.
|
|
newTB := NewWholeTestRun(tb)
|
|
if c.clusters == nil {
|
|
c.clusters = make(map[int]*ShareableCluster)
|
|
// tb should always be a wholeTestRun for clusterCache, and we need to
|
|
// register with that, so our cleanup happens *before* the deletion of
|
|
// the directories, otherwise etcd can fail to flush WAL files on
|
|
// exit, causing tests to fail.
|
|
newTB.Cleanup(c.CleanupClusters)
|
|
}
|
|
cluster, err := underlyingNewCluster(newTB, size)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
cluster.shared = true
|
|
c.clusters[size] = cluster
|
|
return cluster, nil
|
|
}
|
|
|
|
func underlyingNewCluster(tb DirCleaner, size int, opts ...[]server.CommandOption) (*ShareableCluster, error) {
|
|
cluster := &ShareableCluster{Nodes: make([]*Command, size)}
|
|
for i := 0; i < size; i++ {
|
|
var commandOpts []server.CommandOption
|
|
if len(opts) > 0 {
|
|
commandOpts = opts[i%len(opts)]
|
|
}
|
|
m := NewCommandNode(tb, commandOpts...)
|
|
m.Config.ImportWorkerPoolSize = 2
|
|
|
|
cluster.Nodes[i] = m
|
|
}
|
|
// The GetPorts... stuff calls things elsewhere that want a plain testing.TB,
|
|
// and doesn't produce permanent directories, I think.
|
|
err := GetPortsGenConfigs(tb, cluster.Nodes)
|
|
if err != nil {
|
|
return nil, errors.Wrap(err, "configuring cluster ports")
|
|
}
|
|
return cluster, nil
|
|
}
|
|
|
|
var globalClusterCache clusterCache
|
|
|
|
// newCluster creates a new cluster, using the shared cluster cache if no opts are
|
|
// specified.
|
|
func newCluster(tb testing.TB, size int, shareable bool, opts ...[]server.CommandOption) (*Cluster, error) {
|
|
if size == 0 {
|
|
return nil, errors.New("cluster must contain at least one node")
|
|
}
|
|
|
|
if len(opts) != size && len(opts) != 0 && len(opts) != 1 {
|
|
return nil, errors.New("Slice of CommandOptions must be of length 0, 1, or equal to the number of cluster nodes")
|
|
}
|
|
|
|
var shared *ShareableCluster
|
|
var err error
|
|
if !shareable {
|
|
shared, err = underlyingNewCluster(tb, size, opts...)
|
|
} else {
|
|
shared, err = globalClusterCache.newCluster(tb, size)
|
|
}
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return &Cluster{ShareableCluster: shared, tb: tb, indexName: indexName(tb.Name())}, nil
|
|
}
|
|
|
|
// MustRunCluster creates and starts a new cluster. The opts parameter
|
|
// is slightly magical; see MustNewCluster.
|
|
func MustRunCluster(tb testing.TB, size int, opts ...[]server.CommandOption) *Cluster {
|
|
cluster := MustNewCluster(tb, size, opts...)
|
|
if !cluster.started {
|
|
err := cluster.Start()
|
|
if err != nil {
|
|
tb.Fatalf("run cluster: %v", err)
|
|
}
|
|
cluster.started = true
|
|
}
|
|
return cluster
|
|
}
|
|
|
|
// prependOpts applies prependTestServerOpts to each of the ops (one per
|
|
// node, or one for the entire cluser).
|
|
func prependOpts(opts [][]server.CommandOption, size int) [][]server.CommandOption {
|
|
if len(opts) == 0 {
|
|
opts = make([][]server.CommandOption, size)
|
|
for i := 0; i < size; i++ {
|
|
opts[i] = prependTestServerOpts([]server.CommandOption{})
|
|
}
|
|
} else if len(opts) == 1 {
|
|
opts2 := make([][]server.CommandOption, size)
|
|
for i := 0; i < size; i++ {
|
|
opts2[i] = prependTestServerOpts(opts[0])
|
|
}
|
|
return opts2
|
|
} else {
|
|
for i := range opts {
|
|
opts[i] = prependTestServerOpts(opts[i])
|
|
}
|
|
}
|
|
return opts
|
|
}
|
|
|
|
// prependTestServerOpts prepends opts with the OpenInMemTranslateStore,
|
|
// tweaks to initial startup delay, and storage config to disable fsync and
|
|
// specify a smaller RBF size.
|
|
func prependTestServerOpts(opts []server.CommandOption) []server.CommandOption {
|
|
cfg := pilosa.TestHolderConfig()
|
|
defaultOpts := []server.CommandOption{
|
|
server.OptCommandServerOptions(
|
|
pilosa.OptServerOpenTranslateStore(pilosa.OpenInMemTranslateStore),
|
|
pilosa.OptServerNodeDownRetries(5, 100*time.Millisecond),
|
|
pilosa.OptServerStorageConfig(&storage.Config{
|
|
Backend: storage.DefaultBackend,
|
|
FsyncEnabled: false,
|
|
}),
|
|
pilosa.OptServerRBFConfig(cfg.RBFConfig),
|
|
),
|
|
}
|
|
return append(defaultOpts, opts...)
|
|
}
|
|
|
|
func indexName(in string) string {
|
|
return strings.Map(func(r rune) rune {
|
|
if r < 127 {
|
|
switch {
|
|
case unicode.IsLetter(r):
|
|
return unicode.ToLower(r)
|
|
case unicode.IsNumber(r):
|
|
return r
|
|
case r == '/':
|
|
return '_'
|
|
}
|
|
}
|
|
return -1
|
|
}, in)
|
|
}
|