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 }