github.com/m3db/m3@v1.5.0/src/dbnode/integration/client.go (about)

     1  // Copyright (c) 2016 Uber Technologies, Inc.
     2  //
     3  // Permission is hereby granted, free of charge, to any person obtaining a copy
     4  // of this software and associated documentation files (the "Software"), to deal
     5  // in the Software without restriction, including without limitation the rights
     6  // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
     7  // copies of the Software, and to permit persons to whom the Software is
     8  // furnished to do so, subject to the following conditions:
     9  //
    10  // The above copyright notice and this permission notice shall be included in
    11  // all copies or substantial portions of the Software.
    12  //
    13  // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
    14  // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
    15  // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
    16  // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
    17  // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
    18  // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
    19  // THE SOFTWARE.
    20  
    21  package integration
    22  
    23  import (
    24  	"errors"
    25  	"sync"
    26  	"time"
    27  
    28  	"github.com/m3db/m3/src/dbnode/client"
    29  	"github.com/m3db/m3/src/dbnode/generated/thrift/rpc"
    30  	"github.com/m3db/m3/src/dbnode/integration/generate"
    31  	nchannel "github.com/m3db/m3/src/dbnode/network/server/tchannelthrift/node/channel"
    32  	"github.com/m3db/m3/src/dbnode/storage/block"
    33  	"github.com/m3db/m3/src/dbnode/storage/bootstrap/result"
    34  	"github.com/m3db/m3/src/dbnode/topology"
    35  	"github.com/m3db/m3/src/x/checked"
    36  	"github.com/m3db/m3/src/x/ident"
    37  	xsync "github.com/m3db/m3/src/x/sync"
    38  	xtime "github.com/m3db/m3/src/x/time"
    39  
    40  	"github.com/uber/tchannel-go"
    41  	"github.com/uber/tchannel-go/thrift"
    42  )
    43  
    44  // TestTChannelClient is a test only TChannel client that exposes db methods.
    45  type TestTChannelClient struct {
    46  	address string
    47  	channel *tchannel.Channel
    48  	name    string
    49  	node    rpc.TChanNode
    50  }
    51  
    52  // NewTChannelClient creates a new client on the given address.
    53  func NewTChannelClient(name, address string) (*TestTChannelClient, error) {
    54  	channel, err := tchannel.NewChannel(name, nil)
    55  	if err != nil {
    56  		return &TestTChannelClient{}, err
    57  	}
    58  
    59  	endpoint := &thrift.ClientOptions{HostPort: address}
    60  	thriftClient := thrift.NewClient(channel, nchannel.ChannelName, endpoint)
    61  	client := rpc.NewTChanNodeClient(thriftClient)
    62  	return &TestTChannelClient{
    63  		name:    name,
    64  		address: address,
    65  		channel: channel,
    66  		node:    client,
    67  	}, nil
    68  }
    69  
    70  // Address returns the address.
    71  func (client *TestTChannelClient) Address() string {
    72  	return client.address
    73  }
    74  
    75  // Channel returns the TChannel channel.
    76  func (client *TestTChannelClient) Channel() *tchannel.Channel {
    77  	return client.channel
    78  }
    79  
    80  // TChannelClientWrite writes a datapoint using a tchannel client.
    81  func (client *TestTChannelClient) TChannelClientWrite(
    82  	timeout time.Duration, req *rpc.WriteRequest,
    83  ) error {
    84  	ctx, _ := thrift.NewContext(timeout)
    85  	return client.node.Write(ctx, req)
    86  }
    87  
    88  // TChannelClientWriteTagged writes a datapoint using a tchannel client.
    89  func (client *TestTChannelClient) TChannelClientWriteTagged(
    90  	timeout time.Duration, req *rpc.WriteTaggedRequest,
    91  ) error {
    92  	ctx, _ := thrift.NewContext(timeout)
    93  	return client.node.WriteTagged(ctx, req)
    94  }
    95  
    96  // TChannelClientWriteBatch writes a data map using a tchannel client.
    97  func (client *TestTChannelClient) TChannelClientWriteBatch(
    98  	timeout time.Duration, namespace ident.ID, seriesList generate.SeriesBlock,
    99  ) error {
   100  	var elems []*rpc.WriteBatchRawRequestElement
   101  	for _, series := range seriesList {
   102  		for _, dp := range series.Data {
   103  			elem := &rpc.WriteBatchRawRequestElement{
   104  				ID: series.ID.Bytes(),
   105  				Datapoint: &rpc.Datapoint{
   106  					Timestamp:         dp.TimestampNanos.Seconds(),
   107  					Value:             dp.Value,
   108  					Annotation:        dp.Annotation,
   109  					TimestampTimeType: rpc.TimeType_UNIX_SECONDS,
   110  				},
   111  			}
   112  			elems = append(elems, elem)
   113  		}
   114  	}
   115  
   116  	ctx, _ := thrift.NewContext(timeout)
   117  	batchReq := &rpc.WriteBatchRawRequest{
   118  		NameSpace: namespace.Bytes(),
   119  		Elements:  elems,
   120  	}
   121  	return client.node.WriteBatchRaw(ctx, batchReq)
   122  }
   123  
   124  // TChannelClientWriteTaggedBatchRaw fulfills a wwrite tagged batch raw request
   125  // using a tchannel client.
   126  func (client *TestTChannelClient) TChannelClientWriteTaggedBatchRaw(
   127  	timeout time.Duration, req *rpc.WriteTaggedBatchRawRequest,
   128  ) error {
   129  	ctx, _ := thrift.NewContext(timeout)
   130  	return client.node.WriteTaggedBatchRaw(ctx, req)
   131  }
   132  
   133  // TChannelClientFetch fulfills a fetch request using a tchannel client.
   134  func (client *TestTChannelClient) TChannelClientFetch(
   135  	timeout time.Duration, req *rpc.FetchRequest,
   136  ) (*rpc.FetchResult_, error) {
   137  	ctx, _ := thrift.NewContext(timeout)
   138  	return client.node.Fetch(ctx, req)
   139  }
   140  
   141  // TChannelClientFetchTagged fulfills a fetch by tag request using a tchannel client.
   142  func (client *TestTChannelClient) TChannelClientFetchTagged(
   143  	timeout time.Duration, req *rpc.FetchTaggedRequest,
   144  ) (*rpc.FetchTaggedResult_, error) {
   145  	ctx, _ := thrift.NewContext(timeout)
   146  	return client.node.FetchTagged(ctx, req)
   147  }
   148  
   149  // TChannelClientAggregateTiles runs a request for AggregateTiles.
   150  func (client *TestTChannelClient) TChannelClientAggregateTiles(
   151  	timeout time.Duration, req *rpc.AggregateTilesRequest,
   152  ) (*rpc.AggregateTilesResult_, error) {
   153  	ctx, _ := thrift.NewContext(timeout)
   154  	return client.node.AggregateTiles(ctx, req)
   155  }
   156  
   157  // TChannelClientTruncate fulfills a namespace truncation request using a tchannel client.
   158  func (client *TestTChannelClient) TChannelClientTruncate(
   159  	timeout time.Duration, req *rpc.TruncateRequest,
   160  ) (int64, error) {
   161  	ctx, _ := thrift.NewContext(timeout)
   162  	truncated, err := client.node.Truncate(ctx, req)
   163  	if err != nil {
   164  		return 0, err
   165  	}
   166  	return truncated.NumSeries, nil
   167  }
   168  
   169  // TChannelClientHealth fulfills a client health request using a tchannel client.
   170  func (client *TestTChannelClient) TChannelClientHealth(
   171  	timeout time.Duration,
   172  ) (*rpc.NodeHealthResult_, error) {
   173  	ctx, _ := thrift.NewContext(timeout)
   174  	return client.node.Health(ctx)
   175  }
   176  
   177  func m3dbAdminClient(opts client.AdminOptions) (client.AdminClient, error) {
   178  	return client.NewAdminClient(opts)
   179  }
   180  
   181  // m3dbClientWriteBatch writes a data map using an m3db client.
   182  func m3dbClientWriteBatch(client client.Client, workerPool xsync.WorkerPool, namespace ident.ID, seriesList generate.SeriesBlock) error {
   183  	session, err := client.DefaultSession()
   184  	if err != nil {
   185  		return err
   186  	}
   187  
   188  	var (
   189  		errCh = make(chan error, 1)
   190  		wg    sync.WaitGroup
   191  	)
   192  
   193  	for _, series := range seriesList {
   194  		for _, dp := range series.Data {
   195  			wg.Add(1)
   196  			id, d := series.ID, dp
   197  			workerPool.Go(func() {
   198  				defer wg.Done()
   199  
   200  				if err := session.Write(
   201  					namespace, id, d.TimestampNanos, d.Value, xtime.Second, d.Annotation,
   202  				); err != nil {
   203  					select {
   204  					case errCh <- err:
   205  					default:
   206  					}
   207  				}
   208  			})
   209  		}
   210  	}
   211  
   212  	wg.Wait()
   213  	close(errCh)
   214  
   215  	return <-errCh
   216  }
   217  
   218  // m3dbClientFetch fulfills a fetch request using an m3db client.
   219  func m3dbClientFetch(client client.Client, req *rpc.FetchRequest) ([]generate.TestValue, error) {
   220  	session, err := client.DefaultSession()
   221  	if err != nil {
   222  		return nil, err
   223  	}
   224  
   225  	iter, err := session.Fetch(
   226  		ident.StringID(req.NameSpace),
   227  		ident.StringID(req.ID),
   228  		xtime.FromNormalizedTime(req.RangeStart, time.Second),
   229  		xtime.FromNormalizedTime(req.RangeEnd, time.Second),
   230  	)
   231  	if err != nil {
   232  		return nil, err
   233  	}
   234  
   235  	defer iter.Close()
   236  
   237  	var datapoints []generate.TestValue
   238  	for iter.Next() {
   239  		dp, _, annotation := iter.Current()
   240  		datapoints = append(datapoints, generate.TestValue{Datapoint: dp, Annotation: annotation})
   241  	}
   242  	if err := iter.Err(); err != nil {
   243  		return nil, err
   244  	}
   245  	return datapoints, nil
   246  }
   247  
   248  // m3dbClientTruncate fulfills a truncation request using an m3db client.
   249  func m3dbClientTruncate(c client.Client, req *rpc.TruncateRequest) (int64, error) {
   250  	session, err := c.DefaultSession()
   251  	if err != nil {
   252  		return 0, err
   253  	}
   254  
   255  	adminSession, ok := session.(client.AdminSession)
   256  	if !ok {
   257  		return 0, errors.New("unable to get an admin session")
   258  	}
   259  
   260  	return adminSession.Truncate(ident.BinaryID(checked.NewBytes(req.NameSpace, nil)))
   261  }
   262  
   263  func m3dbClientFetchBlocksMetadata(
   264  	c client.AdminClient,
   265  	namespace ident.ID,
   266  	shards []uint32,
   267  	start, end xtime.UnixNano,
   268  	consistencyLevel topology.ReadConsistencyLevel,
   269  ) (map[uint32][]block.ReplicaMetadata, error) {
   270  	session, err := c.DefaultAdminSession()
   271  	if err != nil {
   272  		return nil, err
   273  	}
   274  
   275  	metadatasByShard := make(map[uint32][]block.ReplicaMetadata, 10)
   276  
   277  	// iterate over all shards
   278  	seen := make(map[string]map[xtime.UnixNano]struct{})
   279  	for _, shardID := range shards {
   280  		// clear seen
   281  		for key := range seen {
   282  			delete(seen, key)
   283  		}
   284  
   285  		var metadatas []block.ReplicaMetadata
   286  		iter, err := session.FetchBlocksMetadataFromPeers(namespace,
   287  			shardID, start, end, consistencyLevel, result.NewOptions())
   288  		if err != nil {
   289  			return nil, err
   290  		}
   291  
   292  		for iter.Next() {
   293  			host, blockMetadata := iter.Current()
   294  			idString := blockMetadata.ID.String()
   295  			seenBlocks, ok := seen[idString]
   296  			if !ok {
   297  				seenBlocks = make(map[xtime.UnixNano]struct{})
   298  				seen[idString] = seenBlocks
   299  			}
   300  			if _, ok := seenBlocks[blockMetadata.Start]; ok {
   301  				continue // Already seen
   302  			}
   303  			seenBlocks[blockMetadata.Start] = struct{}{}
   304  			metadatas = append(metadatas, block.ReplicaMetadata{
   305  				Metadata: blockMetadata,
   306  				Host:     host,
   307  			})
   308  		}
   309  		if err := iter.Err(); err != nil {
   310  			return nil, err
   311  		}
   312  
   313  		if metadatas != nil {
   314  			metadatasByShard[shardID] = metadatas
   315  		}
   316  	}
   317  
   318  	return metadatasByShard, nil
   319  }