├── CHANGELOG.md ├── .excludemetalint ├── .excludelint ├── .excludefmt ├── .gitmodules ├── .excludecoverage ├── .gitignore ├── .travis.yml ├── generated ├── proto │ ├── msgpb │ │ └── msg.proto │ ├── topicpb │ │ └── topic.proto │ └── generate.go └── mocks │ └── generate.go ├── README.md ├── glide.yaml ├── .metalinter.json ├── topic ├── consumption_type_test.go ├── options.go ├── consumption_type.go ├── service.go ├── service_test.go ├── topic_mock.go ├── types.go └── topic_test.go ├── producer ├── writer │ ├── metadata.go │ ├── message_pool.go │ ├── router.go │ ├── message_pool_test.go │ ├── router_mock.go │ ├── options_test.go │ ├── message.go │ ├── shard_writer_mock.go │ ├── consumer_service_writer_mock.go │ ├── writer.go │ └── shard_writer.go ├── options.go ├── buffer │ ├── strategy.go │ ├── strategy_test.go │ ├── types.go │ └── options.go ├── config │ ├── producer_test.go │ ├── producer.go │ ├── buffer_test.go │ ├── buffer.go │ ├── writer_test.go │ └── writer.go ├── producer.go ├── ref_counted.go ├── types.go ├── ref_counted_test.go └── producer_mock.go ├── protocol └── proto │ ├── common.go │ ├── benchmark_test.go │ ├── options.go │ ├── config_test.go │ ├── config.go │ ├── encoder.go │ ├── types.go │ ├── decoder.go │ ├── proto_mock.go │ └── roundtrip_test.go ├── consumer ├── pools.go ├── server.go ├── config_test.go ├── server_test.go ├── consumer_mock.go ├── config.go ├── types.go ├── options.go └── consumer.go ├── glide.lock └── Makefile /CHANGELOG.md: -------------------------------------------------------------------------------- 1 | Changelog 2 | ========= 3 | 4 | -------------------------------------------------------------------------------- /.excludemetalint: -------------------------------------------------------------------------------- 1 | vendor/ 2 | generated/ 3 | _mock.go 4 | -------------------------------------------------------------------------------- /.excludelint: -------------------------------------------------------------------------------- 1 | (vendor/) 2 | (generated/) 3 | (_mock.go) 4 | (_string.go) 5 | -------------------------------------------------------------------------------- /.excludefmt: -------------------------------------------------------------------------------- 1 | (vendor/) 2 | (proto/) 3 | (gen-go/) 4 | (.pb.go) 5 | (_mock.go) 6 | -------------------------------------------------------------------------------- /.gitmodules: -------------------------------------------------------------------------------- 1 | [submodule ".ci"] 2 | path = .ci 3 | url = https://github.com/m3db/ci-scripts 4 | -------------------------------------------------------------------------------- /.excludecoverage: -------------------------------------------------------------------------------- 1 | _mock.go 2 | _gen.go 3 | _matcher.go 4 | generated/ 5 | integration/ 6 | vendor/ 7 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | *.out 2 | *.test 3 | *.xml 4 | *.swp 5 | .idea/ 6 | .vscode/ 7 | *.iml 8 | *.ipr 9 | *.iws 10 | *.cov 11 | *.html 12 | *.tmp 13 | test.log 14 | 15 | # glide manages this 16 | vendor/ 17 | 18 | # Build binaries 19 | bin/ 20 | 21 | # Debug binaries 22 | debug 23 | debug.test 24 | 25 | # Test data 26 | test-data 27 | 28 | # Coverage 29 | integration/coverage_imports.go 30 | -------------------------------------------------------------------------------- /.travis.yml: -------------------------------------------------------------------------------- 1 | language: go 2 | go: 3 | - "1.9.x" 4 | - "1.10.x" 5 | install: make install-ci 6 | sudo: false 7 | dist: trusty 8 | env: 9 | global: 10 | - TEST_TIMEOUT_SCALE=20 11 | - PACKAGE=github.com/m3db/m3msg 12 | matrix: 13 | - MAKE_TARGET="test-ci-unit" 14 | - MAKE_TARGET="test-ci-integration" 15 | - MAKE_TARGET="metalint" 16 | script: "make $MAKE_TARGET" 17 | -------------------------------------------------------------------------------- /generated/proto/msgpb/msg.proto: -------------------------------------------------------------------------------- 1 | syntax = "proto3"; 2 | 3 | package msgpb; 4 | 5 | import "github.com/gogo/protobuf/gogoproto/gogo.proto"; 6 | 7 | message Metadata { 8 | uint64 shard = 1; 9 | uint64 id = 2; 10 | } 11 | 12 | message Message { 13 | Metadata metadata = 1 [(gogoproto.nullable) = false]; 14 | bytes value = 2; 15 | } 16 | 17 | message Ack { 18 | repeated Metadata metadata = 1 [(gogoproto.nullable) = false]; 19 | } 20 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # M3msg [![Coverage Status](https://coveralls.io/repos/github/m3db/m3msg/badge.svg?branch=master)](https://coveralls.io/github/m3db/m3msg?branch=master) 2 | 3 | A partitioned message queueing, routing and delivery library designed for very small messages at very high speeds that don't require disk durability. This makes it quite useful for metrics ingestion pipelines. 4 | 5 |
6 | 7 | This project is released under the [Apache License, Version 2.0](LICENSE). 8 | -------------------------------------------------------------------------------- /glide.yaml: -------------------------------------------------------------------------------- 1 | package: github.com/m3db/m3msg 2 | import: 3 | - package: github.com/m3db/m3x 4 | version: 647a4a05078f018242070cbd71bea6ecba964199 5 | - package: github.com/m3db/m3cluster 6 | version: d971450316604c151719f53afbb95539e6dbafbb 7 | - package: github.com/golang/protobuf 8 | version: v1.0.0 9 | - package: github.com/gogo/protobuf 10 | version: v1.0.0 11 | - package: github.com/golang/mock 12 | version: v1.0.1 13 | - package: github.com/fortytw2/leaktest 14 | version: ^1.2.0 15 | -------------------------------------------------------------------------------- /generated/proto/topicpb/topic.proto: -------------------------------------------------------------------------------- 1 | syntax = "proto3"; 2 | package topicpb; 3 | 4 | message Topic { 5 | string name = 1; 6 | uint32 number_of_shards = 2; 7 | repeated ConsumerService consumer_services = 3; 8 | } 9 | 10 | message ConsumerService { 11 | ServiceID service_id = 1; 12 | ConsumptionType consumption_type = 2; 13 | int64 message_ttl_nanos = 3; 14 | } 15 | 16 | message ServiceID { 17 | string name = 1; 18 | string environment = 2; 19 | string zone = 3; 20 | } 21 | 22 | enum ConsumptionType { 23 | UNKNOWN = 0; 24 | SHARED = 1; 25 | REPLICATED = 2; 26 | } 27 | -------------------------------------------------------------------------------- /.metalinter.json: -------------------------------------------------------------------------------- 1 | { 2 | "Linters": { 3 | "badtime": { 4 | "Command": "badtime", 5 | "Pattern": "PATH:LINE:COL:MESSAGE" 6 | }, 7 | "deadcode": { 8 | "Command": "deadcode -tags integration" 9 | }, 10 | "varcheck": { 11 | "Command": "varcheck -tags integration" 12 | }, 13 | "megacheck": { 14 | "Command": "megacheck -tags integration" 15 | } 16 | }, 17 | "Enable": 18 | [ "golint" 19 | , "deadcode" 20 | , "varcheck" 21 | , "structcheck" 22 | , "goconst" 23 | , "ineffassign" 24 | , "unconvert" 25 | , "misspell" 26 | , "unparam" 27 | , "megacheck" 28 | , "badtime" ], 29 | "Deadline": "3m", 30 | "EnableGC": true 31 | } 32 | -------------------------------------------------------------------------------- /generated/proto/generate.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2018 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 | //go:generate sh -c "$GOPATH/src/$PACKAGE/.ci/proto-gen.sh $GOPATH/src/$PACKAGE/generated/proto $GOPATH/src/$PACKAGE/generated/proto" 22 | 23 | package proto 24 | -------------------------------------------------------------------------------- /topic/consumption_type_test.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2018 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 topic 22 | 23 | import ( 24 | "testing" 25 | 26 | "github.com/stretchr/testify/require" 27 | ) 28 | 29 | func TestNewConsumptionType(t *testing.T) { 30 | ct, err := NewConsumptionType("shared") 31 | require.NoError(t, err) 32 | require.Equal(t, Shared, ct) 33 | 34 | ct, err = NewConsumptionType("replicated") 35 | require.NoError(t, err) 36 | require.Equal(t, Replicated, ct) 37 | 38 | ct, err = NewConsumptionType("bad") 39 | require.Error(t, err) 40 | require.Equal(t, Unknown, ct) 41 | } 42 | -------------------------------------------------------------------------------- /producer/writer/metadata.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2018 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 writer 22 | 23 | import "github.com/m3db/m3msg/generated/proto/msgpb" 24 | 25 | // metadata is the metadata for a message. 26 | type metadata struct { 27 | shard uint64 28 | id uint64 29 | } 30 | 31 | func (m metadata) ToProto(pb *msgpb.Metadata) { 32 | pb.Shard = m.shard 33 | pb.Id = m.id 34 | } 35 | 36 | func (m *metadata) FromProto(pb msgpb.Metadata) { 37 | m.shard = pb.Shard 38 | m.id = pb.Id 39 | } 40 | 41 | func newMetadataFromProto(pb msgpb.Metadata) metadata { 42 | var m metadata 43 | m.FromProto(pb) 44 | return m 45 | } 46 | -------------------------------------------------------------------------------- /producer/options.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2018 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 producer 22 | 23 | type producerOptions struct { 24 | buffer Buffer 25 | writer Writer 26 | } 27 | 28 | // NewOptions creates a Options. 29 | func NewOptions() Options { 30 | return &producerOptions{} 31 | } 32 | 33 | func (opts *producerOptions) Buffer() Buffer { 34 | return opts.buffer 35 | } 36 | 37 | func (opts *producerOptions) SetBuffer(value Buffer) Options { 38 | o := *opts 39 | o.buffer = value 40 | return &o 41 | } 42 | 43 | func (opts *producerOptions) Writer() Writer { 44 | return opts.writer 45 | } 46 | 47 | func (opts *producerOptions) SetWriter(value Writer) Options { 48 | o := *opts 49 | o.writer = value 50 | return &o 51 | } 52 | -------------------------------------------------------------------------------- /protocol/proto/common.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2018 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 proto 22 | 23 | import ( 24 | "encoding/binary" 25 | 26 | "github.com/m3db/m3x/pool" 27 | ) 28 | 29 | const ( 30 | sizeEncodingLength = 4 31 | ) 32 | 33 | var sizeEncodeDecoder = binary.BigEndian 34 | 35 | func growDataBufferIfNeeded(buffer []byte, targetSize int, pool pool.BytesPool) []byte { 36 | if len(buffer) >= targetSize { 37 | return buffer 38 | } 39 | if pool != nil { 40 | pool.Put(buffer) 41 | } 42 | return getByteSliceWithLength(targetSize, pool) 43 | } 44 | 45 | func getByteSliceWithLength(targetSize int, pool pool.BytesPool) []byte { 46 | if pool == nil { 47 | return make([]byte, targetSize) 48 | } 49 | b := pool.Get(targetSize) 50 | // Make sure there is enough length in the slice. 51 | return b[:cap(b)] 52 | } 53 | -------------------------------------------------------------------------------- /producer/buffer/strategy.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2018 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 buffer 22 | 23 | import "fmt" 24 | 25 | var ( 26 | validStrategies = []OnFullStrategy{ 27 | ReturnError, 28 | DropOldest, 29 | } 30 | ) 31 | 32 | // UnmarshalYAML unmarshals OnFullStrategy from yaml. 33 | func (t *OnFullStrategy) UnmarshalYAML(unmarshal func(interface{}) error) error { 34 | var str string 35 | if err := unmarshal(&str); err != nil { 36 | return err 37 | } 38 | var validStrings []string 39 | for _, validStrategy := range validStrategies { 40 | validString := string(validStrategy) 41 | if validString == str { 42 | *t = validStrategy 43 | return nil 44 | } 45 | validStrings = append(validStrings, validString) 46 | } 47 | 48 | return fmt.Errorf("invalid strategy %s, valid strategies are: %v", str, validStrings) 49 | } 50 | -------------------------------------------------------------------------------- /consumer/pools.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2018 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 consumer 22 | 23 | import "github.com/m3db/m3x/pool" 24 | 25 | type messagePool struct { 26 | pool.ObjectPool 27 | } 28 | 29 | func newMessagePool(pOpts pool.ObjectPoolOptions) *messagePool { 30 | iOpts := pOpts.InstrumentOptions() 31 | scope := iOpts.MetricsScope() 32 | return &messagePool{ 33 | ObjectPool: pool.NewObjectPool( 34 | pOpts.SetInstrumentOptions( 35 | iOpts.SetMetricsScope( 36 | scope.Tagged(map[string]string{"pool": "message"}), 37 | ), 38 | ), 39 | ), 40 | } 41 | } 42 | 43 | func (p *messagePool) Init() { 44 | p.ObjectPool.Init(func() interface{} { 45 | return newMessage(p) 46 | }) 47 | } 48 | 49 | func (p *messagePool) Get() *message { 50 | return p.ObjectPool.Get().(*message) 51 | } 52 | 53 | func (p *messagePool) Put(m *message) { 54 | p.ObjectPool.Put(m) 55 | } 56 | -------------------------------------------------------------------------------- /producer/config/producer_test.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2018 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 config 22 | 23 | import ( 24 | "testing" 25 | 26 | "github.com/golang/mock/gomock" 27 | "github.com/m3db/m3cluster/client" 28 | "github.com/m3db/m3x/instrument" 29 | 30 | "github.com/stretchr/testify/require" 31 | yaml "gopkg.in/yaml.v2" 32 | ) 33 | 34 | func TestProducerConfiguration(t *testing.T) { 35 | str := ` 36 | buffer: 37 | maxMessageSize: 10 38 | maxBufferSize: 100 39 | writer: 40 | topicName: testTopic 41 | ` 42 | 43 | var cfg ProducerConfiguration 44 | require.NoError(t, yaml.Unmarshal([]byte(str), &cfg)) 45 | 46 | ctrl := gomock.NewController(t) 47 | defer ctrl.Finish() 48 | 49 | cs := client.NewMockClient(ctrl) 50 | cs.EXPECT().Store(gomock.Any()).Return(nil, nil) 51 | cs.EXPECT().Services(gomock.Any()).Return(nil, nil) 52 | 53 | _, err := cfg.newOptions(cs, instrument.NewOptions()) 54 | require.NoError(t, err) 55 | } 56 | -------------------------------------------------------------------------------- /protocol/proto/benchmark_test.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2018 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 proto 22 | 23 | import ( 24 | "bytes" 25 | "testing" 26 | 27 | "github.com/m3db/m3msg/generated/proto/msgpb" 28 | ) 29 | 30 | func BenchmarkBaseEncodeDecodeRoundTrip(b *testing.B) { 31 | r := bytes.NewReader(nil) 32 | encoder := NewEncoder(NewOptions()) 33 | decoder := NewDecoder(r, NewOptions()) 34 | encodeMsg := msgpb.Message{ 35 | Metadata: msgpb.Metadata{}, 36 | Value: make([]byte, 200), 37 | } 38 | decodeMsg := msgpb.Message{} 39 | b.ReportAllocs() 40 | b.ResetTimer() 41 | for n := 0; n < b.N; n++ { 42 | encodeMsg.Metadata.Id = uint64(n) 43 | err := encoder.Encode(&encodeMsg) 44 | if err != nil { 45 | b.FailNow() 46 | } 47 | r.Reset(encoder.Bytes()) 48 | if err := decoder.Decode(&decodeMsg); err != nil { 49 | b.FailNow() 50 | } 51 | if decodeMsg.Metadata.Id != uint64(n) { 52 | b.FailNow() 53 | } 54 | } 55 | } 56 | -------------------------------------------------------------------------------- /protocol/proto/options.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2018 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 proto 22 | 23 | import ( 24 | "github.com/m3db/m3x/pool" 25 | ) 26 | 27 | var ( 28 | defaultMaxMessageSize = 4 * 1024 * 1024 // 4MB. 29 | ) 30 | 31 | // NewOptions creates a new Options. 32 | func NewOptions() Options { 33 | return &options{ 34 | maxMessageSize: defaultMaxMessageSize, 35 | } 36 | } 37 | 38 | type options struct { 39 | maxMessageSize int 40 | bytesPool pool.BytesPool 41 | } 42 | 43 | func (opts *options) MaxMessageSize() int { 44 | return opts.maxMessageSize 45 | } 46 | 47 | func (opts *options) SetMaxMessageSize(value int) Options { 48 | o := *opts 49 | o.maxMessageSize = value 50 | return &o 51 | } 52 | 53 | func (opts *options) BytesPool() pool.BytesPool { 54 | return opts.bytesPool 55 | } 56 | 57 | func (opts *options) SetBytesPool(value pool.BytesPool) Options { 58 | o := *opts 59 | o.bytesPool = value 60 | return &o 61 | } 62 | -------------------------------------------------------------------------------- /protocol/proto/config_test.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2018 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 proto 22 | 23 | import ( 24 | "testing" 25 | 26 | "github.com/m3db/m3x/instrument" 27 | 28 | "github.com/stretchr/testify/require" 29 | yaml "gopkg.in/yaml.v2" 30 | ) 31 | 32 | func TestConfiguration(t *testing.T) { 33 | str := ` 34 | maxMessageSize: 1024 35 | bytesPool: 36 | buckets: 37 | - capacity: 4 38 | count: 60000 39 | - capacity: 1024 40 | count: 30000 41 | watermark: 42 | lowWatermark: 0.01 43 | highWatermark: 0.02 44 | ` 45 | 46 | var cfg Configuration 47 | require.NoError(t, yaml.Unmarshal([]byte(str), &cfg)) 48 | opts := cfg.NewOptions(instrument.NewOptions()) 49 | require.Equal(t, 1024, opts.MaxMessageSize()) 50 | require.NotNil(t, opts.BytesPool()) 51 | b := opts.BytesPool().Get(2) 52 | require.Equal(t, 0, len(b)) 53 | require.Equal(t, 4, cap(b)) 54 | b = opts.BytesPool().Get(200) 55 | require.Equal(t, 0, len(b)) 56 | require.Equal(t, 1024, cap(b)) 57 | } 58 | -------------------------------------------------------------------------------- /producer/buffer/strategy_test.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2018 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 buffer 22 | 23 | import ( 24 | "testing" 25 | 26 | "github.com/stretchr/testify/require" 27 | yaml "gopkg.in/yaml.v2" 28 | ) 29 | 30 | func TestStrategyYamlUnmarshal(t *testing.T) { 31 | var cfg OnFullStrategy 32 | tests := []struct { 33 | bytes []byte 34 | expectErr bool 35 | expectedStrategy OnFullStrategy 36 | }{ 37 | { 38 | bytes: []byte("dropOldest"), 39 | expectErr: false, 40 | expectedStrategy: DropOldest, 41 | }, 42 | { 43 | bytes: []byte("returnError"), 44 | expectErr: false, 45 | expectedStrategy: ReturnError, 46 | }, 47 | { 48 | bytes: []byte("bad"), 49 | expectErr: true, 50 | }, 51 | } 52 | for _, test := range tests { 53 | err := yaml.Unmarshal(test.bytes, &cfg) 54 | if test.expectErr { 55 | require.Error(t, err) 56 | continue 57 | } 58 | require.Equal(t, test.expectedStrategy, cfg) 59 | } 60 | } 61 | -------------------------------------------------------------------------------- /topic/options.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2018 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 topic 22 | 23 | import ( 24 | "github.com/m3db/m3cluster/client" 25 | "github.com/m3db/m3cluster/kv" 26 | ) 27 | 28 | type serviceOptions struct { 29 | configServiceClient client.Client 30 | kvOpts kv.OverrideOptions 31 | } 32 | 33 | // NewServiceOptions returns new ServiceOptions. 34 | func NewServiceOptions() ServiceOptions { 35 | return &serviceOptions{ 36 | kvOpts: kv.NewOverrideOptions(), 37 | } 38 | } 39 | 40 | func (opts *serviceOptions) ConfigService() client.Client { 41 | return opts.configServiceClient 42 | } 43 | 44 | func (opts *serviceOptions) SetConfigService(c client.Client) ServiceOptions { 45 | o := *opts 46 | o.configServiceClient = c 47 | return &o 48 | } 49 | 50 | func (opts *serviceOptions) KVOverrideOptions() kv.OverrideOptions { 51 | return opts.kvOpts 52 | } 53 | 54 | func (opts *serviceOptions) SetKVOverrideOptions(value kv.OverrideOptions) ServiceOptions { 55 | o := *opts 56 | o.kvOpts = value 57 | return &o 58 | } 59 | -------------------------------------------------------------------------------- /protocol/proto/config.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2018 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 proto 22 | 23 | import ( 24 | "github.com/m3db/m3x/instrument" 25 | "github.com/m3db/m3x/pool" 26 | ) 27 | 28 | // Configuration configures an Encoder or a Decoder. 29 | type Configuration struct { 30 | MaxMessageSize *int `yaml:"maxMessageSize"` 31 | BytesPool *pool.BucketizedPoolConfiguration `yaml:"bytesPool"` 32 | } 33 | 34 | // NewOptions creates a new Options. 35 | func (c *Configuration) NewOptions( 36 | iOpts instrument.Options, 37 | ) Options { 38 | var ( 39 | opts = NewOptions() 40 | ) 41 | if c.MaxMessageSize != nil { 42 | opts = opts.SetMaxMessageSize(*c.MaxMessageSize) 43 | } 44 | if c.BytesPool != nil { 45 | scope := iOpts.MetricsScope() 46 | p := pool.NewBytesPool( 47 | c.BytesPool.NewBuckets(), 48 | c.BytesPool.NewObjectPoolOptions(iOpts.SetMetricsScope(scope.Tagged( 49 | map[string]string{"pool": "bytes"}, 50 | ))), 51 | ) 52 | p.Init() 53 | opts = opts.SetBytesPool(p) 54 | } 55 | return opts 56 | } 57 | -------------------------------------------------------------------------------- /producer/producer.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2018 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 producer 22 | 23 | type producer struct { 24 | Buffer 25 | Writer 26 | } 27 | 28 | // NewProducer returns a new producer. 29 | func NewProducer(opts Options) Producer { 30 | return &producer{ 31 | Buffer: opts.Buffer(), 32 | Writer: opts.Writer(), 33 | } 34 | } 35 | 36 | func (p *producer) Init() error { 37 | p.Buffer.Init() 38 | return p.Writer.Init() 39 | } 40 | 41 | func (p *producer) Produce(m Message) error { 42 | rm, err := p.Buffer.Add(m) 43 | if err != nil { 44 | return err 45 | } 46 | return p.Writer.Write(rm) 47 | } 48 | 49 | func (p *producer) Close(ct CloseType) { 50 | // NB: Must close buffer first, it will start returning errors on 51 | // new writes immediately. Then if the close type is to wait for consumption 52 | // it will block until all messages got consumed. If the close type is to drop 53 | // everything, it will drop everything buffered. 54 | p.Buffer.Close(ct) 55 | // Then we can close writer to clean up outstanding go routines. 56 | p.Writer.Close() 57 | } 58 | -------------------------------------------------------------------------------- /consumer/server.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2018 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 consumer 22 | 23 | import ( 24 | "net" 25 | 26 | "github.com/m3db/m3x/server" 27 | ) 28 | 29 | // NewServer creates a new server. 30 | func NewServer(addr string, opts ServerOptions) server.Server { 31 | return server.NewServer( 32 | addr, 33 | NewHandler(opts.ConsumeFn(), opts.ConsumerOptions()), 34 | opts.ServerOptions(), 35 | ) 36 | } 37 | 38 | type handler struct { 39 | opts Options 40 | mPool *messagePool 41 | consumeFn ConsumeFn 42 | 43 | m metrics 44 | } 45 | 46 | // NewHandler creates a new handler. 47 | func NewHandler(consumeFn ConsumeFn, opts Options) server.Handler { 48 | mPool := newMessagePool(opts.MessagePoolOptions()) 49 | mPool.Init() 50 | return &handler{ 51 | consumeFn: consumeFn, 52 | opts: opts, 53 | mPool: mPool, 54 | m: newConsumerMetrics(opts.InstrumentOptions().MetricsScope()), 55 | } 56 | } 57 | 58 | func (h *handler) Handle(conn net.Conn) { 59 | c := newConsumer(conn, h.mPool, h.opts, h.m) 60 | c.Init() 61 | h.consumeFn(c) 62 | } 63 | 64 | func (h *handler) Close() {} 65 | -------------------------------------------------------------------------------- /producer/writer/message_pool.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2018 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 writer 22 | 23 | import "github.com/m3db/m3x/pool" 24 | 25 | // messagePool is the pool for message. 26 | type messagePool interface { 27 | // Init initializes the pool. 28 | Init() 29 | 30 | // Get gets a message from the pool. 31 | Get() *message 32 | 33 | // Put puts a message to the pool. 34 | Put(m *message) 35 | } 36 | 37 | type msgPool struct { 38 | pool.ObjectPool 39 | } 40 | 41 | func newMessagePool(opts pool.ObjectPoolOptions) messagePool { 42 | iOpts := opts.InstrumentOptions() 43 | scope := iOpts.MetricsScope() 44 | p := pool.NewObjectPool( 45 | opts.SetInstrumentOptions( 46 | iOpts.SetMetricsScope( 47 | scope.Tagged(map[string]string{"pool": "message"}), 48 | ), 49 | ), 50 | ) 51 | return &msgPool{p} 52 | } 53 | 54 | func (p *msgPool) Init() { 55 | p.ObjectPool.Init(func() interface{} { 56 | return newMessage() 57 | }) 58 | } 59 | 60 | func (p *msgPool) Get() *message { 61 | return p.ObjectPool.Get().(*message) 62 | } 63 | 64 | func (p *msgPool) Put(m *message) { 65 | p.ObjectPool.Put(m) 66 | } 67 | -------------------------------------------------------------------------------- /producer/writer/router.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2018 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 writer 22 | 23 | import ( 24 | "fmt" 25 | "sync" 26 | ) 27 | 28 | type ackRouter interface { 29 | // Ack acks the metadata. 30 | Ack(ack metadata) error 31 | 32 | // Register registers a message writer. 33 | Register(replicatedShardID uint64, mw messageWriter) 34 | 35 | // Unregister removes a message writer. 36 | Unregister(replicatedShardID uint64) 37 | } 38 | 39 | type router struct { 40 | sync.RWMutex 41 | 42 | messageWriters map[uint64]messageWriter 43 | } 44 | 45 | func newAckRouter(size int) ackRouter { 46 | return &router{ 47 | messageWriters: make(map[uint64]messageWriter, size), 48 | } 49 | } 50 | 51 | func (r *router) Ack(meta metadata) error { 52 | r.RLock() 53 | mw, ok := r.messageWriters[meta.shard] 54 | r.RUnlock() 55 | if !ok { 56 | // Unexpected. 57 | return fmt.Errorf("can't find shard %v", meta.shard) 58 | } 59 | mw.Ack(meta) 60 | return nil 61 | } 62 | 63 | func (r *router) Register(replicatedShardID uint64, mw messageWriter) { 64 | r.Lock() 65 | r.messageWriters[replicatedShardID] = mw 66 | r.Unlock() 67 | } 68 | 69 | func (r *router) Unregister(replicatedShardID uint64) { 70 | r.Lock() 71 | delete(r.messageWriters, replicatedShardID) 72 | r.Unlock() 73 | } 74 | -------------------------------------------------------------------------------- /producer/config/producer.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2018 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 config 22 | 23 | import ( 24 | "github.com/m3db/m3cluster/client" 25 | "github.com/m3db/m3msg/producer" 26 | "github.com/m3db/m3msg/producer/buffer" 27 | "github.com/m3db/m3msg/producer/writer" 28 | "github.com/m3db/m3x/instrument" 29 | ) 30 | 31 | // ProducerConfiguration configs the producer. 32 | type ProducerConfiguration struct { 33 | Buffer BufferConfiguration `yaml:"buffer"` 34 | Writer WriterConfiguration `yaml:"writer"` 35 | } 36 | 37 | func (c *ProducerConfiguration) newOptions( 38 | cs client.Client, 39 | iOpts instrument.Options, 40 | ) (producer.Options, error) { 41 | wOpts, err := c.Writer.NewOptions(cs, iOpts) 42 | if err != nil { 43 | return nil, err 44 | } 45 | b, err := buffer.NewBuffer(c.Buffer.NewOptions(iOpts)) 46 | if err != nil { 47 | return nil, err 48 | } 49 | return producer.NewOptions(). 50 | SetBuffer(b). 51 | SetWriter(writer.NewWriter(wOpts)), nil 52 | } 53 | 54 | // NewProducer creates new producer. 55 | func (c *ProducerConfiguration) NewProducer( 56 | cs client.Client, 57 | iOpts instrument.Options, 58 | ) (producer.Producer, error) { 59 | opts, err := c.newOptions(cs, iOpts) 60 | if err != nil { 61 | return nil, err 62 | } 63 | return producer.NewProducer(opts), nil 64 | } 65 | -------------------------------------------------------------------------------- /consumer/config_test.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2018 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 consumer 22 | 23 | import ( 24 | "testing" 25 | "time" 26 | 27 | "github.com/m3db/m3x/instrument" 28 | 29 | "github.com/stretchr/testify/require" 30 | yaml "gopkg.in/yaml.v2" 31 | ) 32 | 33 | func TestConfiguration(t *testing.T) { 34 | str := ` 35 | messagePool: 36 | size: 5 37 | ackFlushInterval: 100ms 38 | ackBufferSize: 100 39 | connectionWriteBufferSize: 200 40 | connectionReadBufferSize: 300 41 | encoder: 42 | maxMessageSize: 100 43 | bytesPool: 44 | watermark: 45 | low: 0.001 46 | decoder: 47 | maxMessageSize: 200 48 | bytesPool: 49 | watermark: 50 | high: 0.002 51 | ` 52 | 53 | var cfg Configuration 54 | require.NoError(t, yaml.Unmarshal([]byte(str), &cfg)) 55 | 56 | opts := cfg.NewOptions(instrument.NewOptions()) 57 | require.Equal(t, 5, opts.MessagePoolOptions().Size()) 58 | require.Equal(t, 100*time.Millisecond, opts.AckFlushInterval()) 59 | require.Equal(t, 100, opts.AckBufferSize()) 60 | require.Equal(t, 200, opts.ConnectionWriteBufferSize()) 61 | require.Equal(t, 300, opts.ConnectionReadBufferSize()) 62 | require.Equal(t, 100, opts.EncoderOptions().MaxMessageSize()) 63 | require.NotNil(t, opts.EncoderOptions().BytesPool()) 64 | require.Equal(t, 200, opts.DecoderOptions().MaxMessageSize()) 65 | require.NotNil(t, opts.EncoderOptions().BytesPool()) 66 | require.Equal(t, instrument.NewOptions(), opts.InstrumentOptions()) 67 | } 68 | -------------------------------------------------------------------------------- /generated/mocks/generate.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2018 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 | // mockgen rules for generating mocks for exported interfaces (reflection mode). 22 | //go:generate sh -c "mockgen -package=producer $PACKAGE/producer Message,Producer | mockclean -pkg $PACKAGE/producer -out $GOPATH/src/$PACKAGE/producer/producer_mock.go" 23 | //go:generate sh -c "mockgen -package=consumer $PACKAGE/consumer Message | mockclean -pkg $PACKAGE/consumer -out $GOPATH/src/$PACKAGE/consumer/consumer_mock.go" 24 | //go:generate sh -c "mockgen -package=proto $PACKAGE/protocol/proto Encoder,Decoder | mockclean -pkg $PACKAGE/protocol/proto -out $GOPATH/src/$PACKAGE/protocol/proto/proto_mock.go" 25 | //go:generate sh -c "mockgen -package=topic $PACKAGE/topic Service | mockclean -pkg $PACKAGE/topic -out $GOPATH/src/$PACKAGE/topic/topic_mock.go" 26 | 27 | // mockgen rules for generating mocks for unexported interfaces (file mode). 28 | //go:generate sh -c "mockgen -package=writer -destination=$GOPATH/src/$PACKAGE/producer/writer/consumer_service_writer_mock.go -source=$GOPATH/src/$PACKAGE/producer/writer/consumer_service_writer.go" 29 | //go:generate sh -c "mockgen -package=writer -destination=$GOPATH/src/$PACKAGE/producer/writer/router_mock.go -source=$GOPATH/src/$PACKAGE/producer/writer/router.go" 30 | //go:generate sh -c "mockgen -package=writer -destination=$GOPATH/src/$PACKAGE/producer/writer/shard_writer_mock.go -source=$GOPATH/src/$PACKAGE/producer/writer/shard_writer.go" 31 | 32 | package mocks 33 | -------------------------------------------------------------------------------- /protocol/proto/encoder.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2018 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 proto 22 | 23 | import ( 24 | "fmt" 25 | 26 | "github.com/m3db/m3x/pool" 27 | ) 28 | 29 | type encoder struct { 30 | buffer []byte 31 | bytesPool pool.BytesPool 32 | maxMessageSize int 33 | encoded int 34 | } 35 | 36 | // NewEncoder creates a new encoder, the implementation is not thread safe. 37 | func NewEncoder(opts Options) Encoder { 38 | if opts == nil { 39 | opts = NewOptions() 40 | } 41 | pool := opts.BytesPool() 42 | return &encoder{ 43 | buffer: getByteSliceWithLength(sizeEncodingLength, pool), 44 | bytesPool: pool, 45 | maxMessageSize: opts.MaxMessageSize(), 46 | } 47 | } 48 | 49 | func (e *encoder) Encode(m Marshaler) error { 50 | size := m.Size() 51 | if size > e.maxMessageSize { 52 | return fmt.Errorf("message size %d is larger than maximum supported size %d", size, e.maxMessageSize) 53 | } 54 | e.buffer = growDataBufferIfNeeded(e.buffer, sizeEncodingLength+size, e.bytesPool) 55 | e.encodeSize(size) 56 | if err := e.encodeData(e.buffer[sizeEncodingLength:], m); err != nil { 57 | return err 58 | } 59 | e.encoded = sizeEncodingLength + size 60 | return nil 61 | } 62 | 63 | func (e *encoder) Bytes() []byte { 64 | return e.buffer[:e.encoded] 65 | } 66 | 67 | func (e *encoder) encodeSize(size int) { 68 | sizeEncodeDecoder.PutUint32(e.buffer, uint32(size)) 69 | } 70 | 71 | func (e *encoder) encodeData(buffer []byte, m Marshaler) error { 72 | _, err := m.MarshalTo(buffer) 73 | return err 74 | } 75 | -------------------------------------------------------------------------------- /protocol/proto/types.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2018 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 proto 22 | 23 | import ( 24 | "io" 25 | 26 | "github.com/m3db/m3x/pool" 27 | ) 28 | 29 | // Marshaler can be marshaled. 30 | type Marshaler interface { 31 | // Size returns the size of the marshaled bytes. 32 | Size() int 33 | 34 | // MarshalTo marshals the marshaler into the given byte slice. 35 | MarshalTo(data []byte) (int, error) 36 | } 37 | 38 | // Unmarshaler can be unmarshaled from bytes. 39 | type Unmarshaler interface { 40 | // Unmarshal unmarshals the unmarshaler from the given byte slice. 41 | Unmarshal(data []byte) error 42 | } 43 | 44 | // Encoder encodes the marshaler. 45 | type Encoder interface { 46 | // Encode encodes the marshaler. 47 | Encode(m Marshaler) error 48 | 49 | // Bytes returns the encoded bytes, the bytes could be reused by 50 | // the next encode call. 51 | Bytes() []byte 52 | } 53 | 54 | // Decoder decodes into an unmarshaler. 55 | type Decoder interface { 56 | // Decode decodes the unmarshaler. 57 | Decode(m Unmarshaler) error 58 | 59 | // ResetReader resets the reader. 60 | ResetReader(r io.Reader) 61 | } 62 | 63 | // Options configures a encoder or decoder. 64 | type Options interface { 65 | // MaxMessageSize returns the maximum message size. 66 | MaxMessageSize() int 67 | 68 | // SetMaxMessageSize sets the maximum message size. 69 | SetMaxMessageSize(value int) Options 70 | 71 | // BytesPool returns the bytes pool. 72 | BytesPool() pool.BytesPool 73 | 74 | // SetBytesPool sets the bytes pool. 75 | SetBytesPool(value pool.BytesPool) Options 76 | } 77 | -------------------------------------------------------------------------------- /producer/writer/message_pool_test.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2018 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 writer 22 | 23 | import ( 24 | "testing" 25 | 26 | "github.com/m3db/m3msg/generated/proto/msgpb" 27 | "github.com/m3db/m3msg/producer" 28 | "github.com/m3db/m3x/pool" 29 | 30 | "github.com/golang/mock/gomock" 31 | "github.com/stretchr/testify/require" 32 | ) 33 | 34 | func TestMessagePool(t *testing.T) { 35 | p := newMessagePool(pool.NewObjectPoolOptions().SetSize(1)) 36 | p.Init() 37 | 38 | ctrl := gomock.NewController(t) 39 | defer ctrl.Finish() 40 | 41 | mm := producer.NewMockMessage(ctrl) 42 | mm.EXPECT().Size().Return(3) 43 | rm := producer.NewRefCountedMessage(mm, nil) 44 | rm.IncRef() 45 | 46 | m := p.Get() 47 | require.Nil(t, m.pb.Value) 48 | mm.EXPECT().Bytes().Return([]byte("foo")) 49 | m.Set(metadata{}, rm, 500) 50 | m.SetRetryAtNanos(100) 51 | 52 | pb, ok := m.Marshaler() 53 | require.True(t, ok) 54 | require.Equal(t, []byte("foo"), pb.(*msgpb.Message).Value) 55 | 56 | mm.EXPECT().Finalize(producer.Consumed) 57 | m.Ack() 58 | require.True(t, m.IsDroppedOrConsumed()) 59 | require.NotNil(t, m.pb.Value) 60 | m.Close() 61 | require.Nil(t, m.pb.Value) 62 | p.Put(m) 63 | 64 | m = p.Get() 65 | require.Nil(t, m.pb.Value) 66 | require.True(t, m.IsDroppedOrConsumed()) 67 | require.Equal(t, int64(500), m.InitNanos()) 68 | 69 | mm.EXPECT().Size().Return(3) 70 | mm.EXPECT().Bytes().Return([]byte("foo")) 71 | m.Set(metadata{}, producer.NewRefCountedMessage(mm, nil), 600) 72 | require.False(t, m.IsDroppedOrConsumed()) 73 | require.Equal(t, int64(600), m.InitNanos()) 74 | } 75 | -------------------------------------------------------------------------------- /consumer/server_test.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2018 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 consumer 22 | 23 | import ( 24 | "net" 25 | "sync" 26 | "testing" 27 | 28 | "github.com/m3db/m3msg/generated/proto/msgpb" 29 | "github.com/m3db/m3msg/protocol/proto" 30 | 31 | "github.com/fortytw2/leaktest" 32 | "github.com/stretchr/testify/require" 33 | ) 34 | 35 | func TestConsumerServer(t *testing.T) { 36 | defer leaktest.Check(t)() 37 | 38 | var ( 39 | count = 0 40 | bytes []byte 41 | closed bool 42 | wg sync.WaitGroup 43 | ) 44 | consumeFn := func(c Consumer) { 45 | for { 46 | count++ 47 | m, err := c.Message() 48 | if err != nil { 49 | break 50 | } 51 | bytes = m.Bytes() 52 | m.Ack() 53 | wg.Done() 54 | } 55 | c.Close() 56 | closed = true 57 | } 58 | 59 | // Set a large ack buffer size to make sure the background go routine 60 | // can flush it. 61 | opts := NewServerOptions().SetConsumerOptions(testOptions().SetAckBufferSize(100)).SetConsumeFn(consumeFn) 62 | l, err := net.Listen("tcp", "127.0.0.1:0") 63 | require.NoError(t, err) 64 | 65 | s := NewServer("a", opts) 66 | s.Serve(l) 67 | 68 | conn, err := net.Dial("tcp", l.Addr().String()) 69 | require.NoError(t, err) 70 | 71 | wg.Add(1) 72 | err = produce(conn, &testMsg1) 73 | require.NoError(t, err) 74 | 75 | wg.Wait() 76 | require.Equal(t, testMsg1.Value, bytes) 77 | 78 | var ack msgpb.Ack 79 | testDecoder := proto.NewDecoder(conn, opts.ConsumerOptions().DecoderOptions()) 80 | err = testDecoder.Decode(&ack) 81 | require.NoError(t, err) 82 | require.Equal(t, 1, len(ack.Metadata)) 83 | require.Equal(t, testMsg1.Metadata, ack.Metadata[0]) 84 | 85 | s.Close() 86 | require.True(t, closed) 87 | } 88 | -------------------------------------------------------------------------------- /producer/config/buffer_test.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2018 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 config 22 | 23 | import ( 24 | "testing" 25 | "time" 26 | 27 | "github.com/m3db/m3msg/producer/buffer" 28 | "github.com/m3db/m3x/instrument" 29 | "github.com/m3db/m3x/retry" 30 | 31 | "github.com/stretchr/testify/require" 32 | "gopkg.in/yaml.v2" 33 | ) 34 | 35 | func TestBufferConfiguration(t *testing.T) { 36 | str := ` 37 | onFullStrategy: returnError 38 | maxBufferSize: 100 39 | maxMessageSize: 16 40 | closeCheckInterval: 3s 41 | scanBatchSize: 128 42 | dropOldestInterval: 500ms 43 | allowedSpilloverRatio: 0.1 44 | cleanupRetry: 45 | initialBackoff: 2s 46 | ` 47 | 48 | var cfg BufferConfiguration 49 | require.NoError(t, yaml.Unmarshal([]byte(str), &cfg)) 50 | 51 | bOpts := cfg.NewOptions(instrument.NewOptions()) 52 | require.Equal(t, buffer.ReturnError, bOpts.OnFullStrategy()) 53 | require.Equal(t, 100, bOpts.MaxBufferSize()) 54 | require.Equal(t, 16, bOpts.MaxMessageSize()) 55 | require.Equal(t, 3*time.Second, bOpts.CloseCheckInterval()) 56 | require.Equal(t, 128, bOpts.ScanBatchSize()) 57 | require.Equal(t, 500*time.Millisecond, bOpts.DropOldestInterval()) 58 | require.Equal(t, 0.1, bOpts.AllowedSpilloverRatio()) 59 | require.Equal(t, 2*time.Second, bOpts.CleanupRetryOptions().InitialBackoff()) 60 | } 61 | 62 | func TestEmptyBufferConfiguration(t *testing.T) { 63 | var cfg BufferConfiguration 64 | require.NoError(t, yaml.Unmarshal(nil, &cfg)) 65 | require.Equal(t, BufferConfiguration{}, cfg) 66 | rOpts := retry.NewOptions() 67 | require.Equal(t, 68 | buffer.NewOptions().SetCleanupRetryOptions(rOpts), 69 | cfg.NewOptions(instrument.NewOptions()).SetCleanupRetryOptions(rOpts), 70 | ) 71 | } 72 | -------------------------------------------------------------------------------- /topic/consumption_type.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2018 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 topic 22 | 23 | import ( 24 | "fmt" 25 | "strings" 26 | 27 | "github.com/m3db/m3msg/generated/proto/topicpb" 28 | ) 29 | 30 | var ( 31 | validTypes = []ConsumptionType{ 32 | Shared, 33 | Replicated, 34 | } 35 | ) 36 | 37 | func (t ConsumptionType) String() string { 38 | return string(t) 39 | } 40 | 41 | // NewConsumptionType creates a consumption type from a string. 42 | func NewConsumptionType(str string) (ConsumptionType, error) { 43 | var validTypeStrings []string 44 | for _, t := range validTypes { 45 | validStr := t.String() 46 | if validStr == str { 47 | return t, nil 48 | } 49 | validTypeStrings = append(validTypeStrings, validStr) 50 | } 51 | return Unknown, fmt.Errorf( 52 | "invalid consumption type '%s', valid types are: %s", str, strings.Join(validTypeStrings, ", "), 53 | ) 54 | } 55 | 56 | // NewConsumptionTypeFromProto creates ConsumptionType from a proto. 57 | func NewConsumptionTypeFromProto(ct topicpb.ConsumptionType) (ConsumptionType, error) { 58 | switch ct { 59 | case topicpb.ConsumptionType_SHARED: 60 | return Shared, nil 61 | case topicpb.ConsumptionType_REPLICATED: 62 | return Replicated, nil 63 | } 64 | return Unknown, fmt.Errorf("invalid consumption type in protobuf: %v", ct) 65 | } 66 | 67 | // ConsumptionTypeToProto creates proto from a ConsumptionType. 68 | func ConsumptionTypeToProto(ct ConsumptionType) (topicpb.ConsumptionType, error) { 69 | switch ct { 70 | case Shared: 71 | return topicpb.ConsumptionType_SHARED, nil 72 | case Replicated: 73 | return topicpb.ConsumptionType_REPLICATED, nil 74 | } 75 | return topicpb.ConsumptionType_UNKNOWN, fmt.Errorf("invalid consumption type: %v", ct) 76 | } 77 | -------------------------------------------------------------------------------- /protocol/proto/decoder.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2018 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 proto 22 | 23 | import ( 24 | "fmt" 25 | "io" 26 | 27 | "github.com/m3db/m3x/pool" 28 | ) 29 | 30 | type decoder struct { 31 | r io.Reader 32 | buffer []byte 33 | bytesPool pool.BytesPool 34 | maxMessageSize int 35 | } 36 | 37 | // NewDecoder decodes a new decoder, the implementation is not thread safe. 38 | func NewDecoder(r io.Reader, opts Options) Decoder { 39 | if opts == nil { 40 | opts = NewOptions() 41 | } 42 | pool := opts.BytesPool() 43 | return &decoder{ 44 | r: r, 45 | buffer: getByteSliceWithLength(sizeEncodingLength, pool), 46 | bytesPool: pool, 47 | maxMessageSize: opts.MaxMessageSize(), 48 | } 49 | } 50 | 51 | func (d *decoder) Decode(m Unmarshaler) error { 52 | size, err := d.decodeSize() 53 | if err != nil { 54 | return err 55 | } 56 | d.buffer = growDataBufferIfNeeded(d.buffer, sizeEncodingLength+size, d.bytesPool) 57 | if size > d.maxMessageSize { 58 | return fmt.Errorf("decoded message size %d is larger than maximum supported size %d", size, d.maxMessageSize) 59 | } 60 | return d.decodeData(d.buffer[sizeEncodingLength:sizeEncodingLength+size], m) 61 | } 62 | 63 | func (d *decoder) decodeSize() (int, error) { 64 | if _, err := io.ReadFull(d.r, d.buffer[:sizeEncodingLength]); err != nil { 65 | return 0, err 66 | } 67 | size := sizeEncodeDecoder.Uint32(d.buffer[:sizeEncodingLength]) 68 | return int(size), nil 69 | } 70 | 71 | func (d *decoder) decodeData(buffer []byte, m Unmarshaler) error { 72 | if _, err := io.ReadFull(d.r, buffer); err != nil { 73 | return err 74 | } 75 | return m.Unmarshal(buffer) 76 | } 77 | 78 | func (d *decoder) ResetReader(r io.Reader) { 79 | d.r = r 80 | } 81 | -------------------------------------------------------------------------------- /consumer/consumer_mock.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2018 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 | // Code generated by MockGen. DO NOT EDIT. 22 | // Source: github.com/m3db/m3msg/consumer (interfaces: Message) 23 | 24 | // Package consumer is a generated GoMock package. 25 | package consumer 26 | 27 | import ( 28 | "reflect" 29 | 30 | "github.com/golang/mock/gomock" 31 | ) 32 | 33 | // MockMessage is a mock of Message interface 34 | type MockMessage struct { 35 | ctrl *gomock.Controller 36 | recorder *MockMessageMockRecorder 37 | } 38 | 39 | // MockMessageMockRecorder is the mock recorder for MockMessage 40 | type MockMessageMockRecorder struct { 41 | mock *MockMessage 42 | } 43 | 44 | // NewMockMessage creates a new mock instance 45 | func NewMockMessage(ctrl *gomock.Controller) *MockMessage { 46 | mock := &MockMessage{ctrl: ctrl} 47 | mock.recorder = &MockMessageMockRecorder{mock} 48 | return mock 49 | } 50 | 51 | // EXPECT returns an object that allows the caller to indicate expected use 52 | func (m *MockMessage) EXPECT() *MockMessageMockRecorder { 53 | return m.recorder 54 | } 55 | 56 | // Ack mocks base method 57 | func (m *MockMessage) Ack() { 58 | m.ctrl.Call(m, "Ack") 59 | } 60 | 61 | // Ack indicates an expected call of Ack 62 | func (mr *MockMessageMockRecorder) Ack() *gomock.Call { 63 | return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Ack", reflect.TypeOf((*MockMessage)(nil).Ack)) 64 | } 65 | 66 | // Bytes mocks base method 67 | func (m *MockMessage) Bytes() []byte { 68 | ret := m.ctrl.Call(m, "Bytes") 69 | ret0, _ := ret[0].([]byte) 70 | return ret0 71 | } 72 | 73 | // Bytes indicates an expected call of Bytes 74 | func (mr *MockMessageMockRecorder) Bytes() *gomock.Call { 75 | return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Bytes", reflect.TypeOf((*MockMessage)(nil).Bytes)) 76 | } 77 | -------------------------------------------------------------------------------- /consumer/config.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2018 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 consumer 22 | 23 | import ( 24 | "time" 25 | 26 | "github.com/m3db/m3msg/protocol/proto" 27 | "github.com/m3db/m3x/instrument" 28 | "github.com/m3db/m3x/pool" 29 | ) 30 | 31 | // Configuration configs the consumer options. 32 | type Configuration struct { 33 | Encoder *proto.Configuration `yaml:"encoder"` 34 | Decoder *proto.Configuration `yaml:"decoder"` 35 | MessagePool *pool.ObjectPoolConfiguration `yaml:"messagePool"` 36 | AckFlushInterval *time.Duration `yaml:"ackFlushInterval"` 37 | AckBufferSize *int `yaml:"ackBufferSize"` 38 | ConnectionWriteBufferSize *int `yaml:"connectionWriteBufferSize"` 39 | ConnectionReadBufferSize *int `yaml:"connectionReadBufferSize"` 40 | } 41 | 42 | // NewOptions creates consumer options. 43 | func (c *Configuration) NewOptions(iOpts instrument.Options) Options { 44 | opts := NewOptions().SetInstrumentOptions(iOpts) 45 | if c.Encoder != nil { 46 | opts = opts.SetEncoderOptions(c.Encoder.NewOptions(iOpts)) 47 | } 48 | if c.Decoder != nil { 49 | opts = opts.SetDecoderOptions(c.Decoder.NewOptions(iOpts)) 50 | } 51 | if c.MessagePool != nil { 52 | opts = opts.SetMessagePoolOptions(c.MessagePool.NewObjectPoolOptions(iOpts)) 53 | } 54 | if c.AckFlushInterval != nil { 55 | opts = opts.SetAckFlushInterval(*c.AckFlushInterval) 56 | } 57 | if c.AckBufferSize != nil { 58 | opts = opts.SetAckBufferSize(*c.AckBufferSize) 59 | } 60 | if c.ConnectionWriteBufferSize != nil { 61 | opts = opts.SetConnectionWriteBufferSize(*c.ConnectionWriteBufferSize) 62 | } 63 | if c.ConnectionReadBufferSize != nil { 64 | opts = opts.SetConnectionReadBufferSize(*c.ConnectionReadBufferSize) 65 | } 66 | return opts 67 | } 68 | -------------------------------------------------------------------------------- /producer/config/buffer.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2018 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 config 22 | 23 | import ( 24 | "time" 25 | 26 | "github.com/m3db/m3msg/producer/buffer" 27 | "github.com/m3db/m3x/instrument" 28 | "github.com/m3db/m3x/retry" 29 | ) 30 | 31 | // BufferConfiguration configs the buffer. 32 | type BufferConfiguration struct { 33 | OnFullStrategy *buffer.OnFullStrategy `yaml:"onFullStrategy"` 34 | MaxBufferSize *int `yaml:"maxBufferSize"` 35 | MaxMessageSize *int `yaml:"maxMessageSize"` 36 | CloseCheckInterval *time.Duration `yaml:"closeCheckInterval"` 37 | DropOldestInterval *time.Duration `yaml:"dropOldestInterval"` 38 | ScanBatchSize *int `yaml:"scanBatchSize"` 39 | AllowedSpilloverRatio *float64 `yaml:"allowedSpilloverRatio"` 40 | CleanupRetry *retry.Configuration `yaml:"cleanupRetry"` 41 | } 42 | 43 | // NewOptions creates new buffer options. 44 | func (c *BufferConfiguration) NewOptions(iOpts instrument.Options) buffer.Options { 45 | opts := buffer.NewOptions().SetOnFullStrategy(buffer.DropOldest) 46 | if c.MaxBufferSize != nil { 47 | opts = opts.SetMaxBufferSize(*c.MaxBufferSize) 48 | } 49 | if c.MaxMessageSize != nil { 50 | opts = opts.SetMaxMessageSize(*c.MaxMessageSize) 51 | } 52 | if c.CloseCheckInterval != nil { 53 | opts = opts.SetCloseCheckInterval(*c.CloseCheckInterval) 54 | } 55 | if c.OnFullStrategy != nil { 56 | opts = opts.SetOnFullStrategy(*c.OnFullStrategy) 57 | } 58 | if c.DropOldestInterval != nil { 59 | opts = opts.SetDropOldestInterval(*c.DropOldestInterval) 60 | } 61 | if c.ScanBatchSize != nil { 62 | opts = opts.SetScanBatchSize(*c.ScanBatchSize) 63 | } 64 | if c.AllowedSpilloverRatio != nil { 65 | opts = opts.SetAllowedSpilloverRatio(*c.AllowedSpilloverRatio) 66 | } 67 | if c.CleanupRetry != nil { 68 | opts = opts.SetCleanupRetryOptions(c.CleanupRetry.NewOptions(iOpts.MetricsScope())) 69 | } 70 | return opts.SetInstrumentOptions(iOpts) 71 | } 72 | -------------------------------------------------------------------------------- /producer/writer/router_mock.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2018 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 | // Code generated by MockGen. DO NOT EDIT. 22 | // Source: github.com/m3db/m3msg/producer/writer/router.go 23 | 24 | // Package writer is a generated GoMock package. 25 | package writer 26 | 27 | import ( 28 | "reflect" 29 | 30 | "github.com/golang/mock/gomock" 31 | ) 32 | 33 | // MockackRouter is a mock of ackRouter interface 34 | type MockackRouter struct { 35 | ctrl *gomock.Controller 36 | recorder *MockackRouterMockRecorder 37 | } 38 | 39 | // MockackRouterMockRecorder is the mock recorder for MockackRouter 40 | type MockackRouterMockRecorder struct { 41 | mock *MockackRouter 42 | } 43 | 44 | // NewMockackRouter creates a new mock instance 45 | func NewMockackRouter(ctrl *gomock.Controller) *MockackRouter { 46 | mock := &MockackRouter{ctrl: ctrl} 47 | mock.recorder = &MockackRouterMockRecorder{mock} 48 | return mock 49 | } 50 | 51 | // EXPECT returns an object that allows the caller to indicate expected use 52 | func (m *MockackRouter) EXPECT() *MockackRouterMockRecorder { 53 | return m.recorder 54 | } 55 | 56 | // Ack mocks base method 57 | func (m *MockackRouter) Ack(ack metadata) error { 58 | ret := m.ctrl.Call(m, "Ack", ack) 59 | ret0, _ := ret[0].(error) 60 | return ret0 61 | } 62 | 63 | // Ack indicates an expected call of Ack 64 | func (mr *MockackRouterMockRecorder) Ack(ack interface{}) *gomock.Call { 65 | return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Ack", reflect.TypeOf((*MockackRouter)(nil).Ack), ack) 66 | } 67 | 68 | // Register mocks base method 69 | func (m *MockackRouter) Register(replicatedShardID uint64, mw messageWriter) { 70 | m.ctrl.Call(m, "Register", replicatedShardID, mw) 71 | } 72 | 73 | // Register indicates an expected call of Register 74 | func (mr *MockackRouterMockRecorder) Register(replicatedShardID, mw interface{}) *gomock.Call { 75 | return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Register", reflect.TypeOf((*MockackRouter)(nil).Register), replicatedShardID, mw) 76 | } 77 | 78 | // Unregister mocks base method 79 | func (m *MockackRouter) Unregister(replicatedShardID uint64) { 80 | m.ctrl.Call(m, "Unregister", replicatedShardID) 81 | } 82 | 83 | // Unregister indicates an expected call of Unregister 84 | func (mr *MockackRouterMockRecorder) Unregister(replicatedShardID interface{}) *gomock.Call { 85 | return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Unregister", reflect.TypeOf((*MockackRouter)(nil).Unregister), replicatedShardID) 86 | } 87 | -------------------------------------------------------------------------------- /topic/service.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2018 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 topic 22 | 23 | import ( 24 | "errors" 25 | 26 | "github.com/m3db/m3cluster/kv" 27 | ) 28 | 29 | var ( 30 | defaultNamespace = "/topic" 31 | errTopicNotAvailable = errors.New("topic is not available") 32 | ) 33 | 34 | type service struct { 35 | store kv.Store 36 | } 37 | 38 | // NewService creates a topic service. 39 | func NewService(sOpts ServiceOptions) (Service, error) { 40 | kvOpts := sanitizeKVOptions(sOpts.KVOverrideOptions()) 41 | store, err := sOpts.ConfigService().Store(kvOpts) 42 | if err != nil { 43 | return nil, err 44 | } 45 | return &service{ 46 | store: store, 47 | }, nil 48 | } 49 | 50 | func (s *service) Get(name string) (Topic, error) { 51 | value, err := s.store.Get(key(name)) 52 | if err != nil { 53 | return nil, err 54 | } 55 | t, err := NewTopicFromValue(value) 56 | if err != nil { 57 | return nil, err 58 | } 59 | return t, nil 60 | } 61 | 62 | func (s *service) CheckAndSet(t Topic, version int) (Topic, error) { 63 | if err := t.Validate(); err != nil { 64 | return nil, err 65 | } 66 | pb, err := ToProto(t) 67 | if err != nil { 68 | return nil, err 69 | } 70 | version, err = s.store.CheckAndSet(key(t.Name()), version, pb) 71 | if err != nil { 72 | return nil, err 73 | } 74 | return t.SetVersion(version), nil 75 | } 76 | 77 | func (s *service) Delete(name string) error { 78 | _, err := s.store.Delete(key(name)) 79 | return err 80 | } 81 | 82 | func (s *service) Watch(name string) (Watch, error) { 83 | w, err := s.store.Watch(key(name)) 84 | if err != nil { 85 | return nil, err 86 | } 87 | return NewWatch(w), nil 88 | } 89 | 90 | func key(name string) string { 91 | return name 92 | } 93 | 94 | func sanitizeKVOptions(opts kv.OverrideOptions) kv.OverrideOptions { 95 | if opts.Namespace() == "" { 96 | opts = opts.SetNamespace(defaultNamespace) 97 | } 98 | return opts 99 | } 100 | 101 | // NewWatch creates a new topic watch. 102 | func NewWatch(w kv.ValueWatch) Watch { 103 | return &watch{w} 104 | } 105 | 106 | type watch struct { 107 | kv.ValueWatch 108 | } 109 | 110 | func (w *watch) Get() (Topic, error) { 111 | value := w.ValueWatch.Get() 112 | if value == nil { 113 | return nil, errTopicNotAvailable 114 | } 115 | t, err := NewTopicFromValue(value) 116 | if err != nil { 117 | return nil, err 118 | } 119 | return t, nil 120 | } 121 | -------------------------------------------------------------------------------- /glide.lock: -------------------------------------------------------------------------------- 1 | hash: 162187b91d0d35f8e950cff7bc8fb398d6d77f5231a82fbd6cdbd2ad55da1156 2 | updated: 2018-07-19T12:14:46.451346-04:00 3 | imports: 4 | - name: github.com/apache/thrift 5 | version: b2a4d4ae21c789b689dd162deb819665567f481c 6 | subpackages: 7 | - lib/go/thrift 8 | - name: github.com/beorn7/perks 9 | version: 3ac7bf7a47d159a033b107610db8a1b6575507a4 10 | subpackages: 11 | - quantile 12 | - name: github.com/fortytw2/leaktest 13 | version: a5ef70473c97b71626b9abeda80ee92ba2a7de9e 14 | - name: github.com/gogo/protobuf 15 | version: 1adfc126b41513cc696b209667c8656ea7aac67c 16 | subpackages: 17 | - gogoproto 18 | - proto 19 | - protoc-gen-gogo/descriptor 20 | - name: github.com/golang/mock 21 | version: 0d6884b72b943387bb12268d5178dde0e759e49a 22 | subpackages: 23 | - gomock 24 | - name: github.com/golang/protobuf 25 | version: 925541529c1fa6821df4e44ce2723319eb2be768 26 | subpackages: 27 | - proto 28 | - name: github.com/google/go-cmp 29 | version: 3af367b6b30c263d47e8895973edcca9a49cf029 30 | subpackages: 31 | - cmp 32 | - cmp/internal/diff 33 | - cmp/internal/function 34 | - cmp/internal/value 35 | - name: github.com/m3db/m3cluster 36 | version: d971450316604c151719f53afbb95539e6dbafbb 37 | subpackages: 38 | - client 39 | - generated/proto/metadatapb 40 | - generated/proto/placementpb 41 | - kv 42 | - kv/mem 43 | - kv/util/runtime 44 | - placement 45 | - placement/algo 46 | - placement/selector 47 | - placement/service 48 | - placement/storage 49 | - services 50 | - services/leader/campaign 51 | - shard 52 | - name: github.com/m3db/m3x 53 | version: 647a4a05078f018242070cbd71bea6ecba964199 54 | subpackages: 55 | - checked 56 | - clock 57 | - close 58 | - errors 59 | - instrument 60 | - log 61 | - net 62 | - pool 63 | - process 64 | - resource 65 | - retry 66 | - server 67 | - sync 68 | - test 69 | - watch 70 | - name: github.com/m3db/prometheus_client_golang 71 | version: 8ae269d24972b8695572fa6b2e3718b5ea82d6b4 72 | subpackages: 73 | - prometheus 74 | - prometheus/promhttp 75 | - name: github.com/m3db/prometheus_client_model 76 | version: 8b2299a4bf7d7fc10835527021716d4b4a6e8700 77 | subpackages: 78 | - go 79 | - name: github.com/m3db/prometheus_common 80 | version: 25aaa3dff79bb48116615ebe1dea6a494b74ce77 81 | subpackages: 82 | - expfmt 83 | - internal/bitbucket.org/ww/goautoneg 84 | - model 85 | - name: github.com/m3db/prometheus_procfs 86 | version: 1878d9fbb537119d24b21ca07effd591627cd160 87 | - name: github.com/matttproud/golang_protobuf_extensions 88 | version: fc2b8d3a73c4867e51861bbdd5ae3c1f0869dd6a 89 | subpackages: 90 | - pbutil 91 | - name: github.com/stretchr/testify 92 | version: efa3c4c36479edac5af79cbcb4bc8cb525e09b13 93 | subpackages: 94 | - assert 95 | - require 96 | - name: github.com/uber-go/tally 97 | version: ff17f3c43c065c3c2991f571e740eee43ea3a14a 98 | subpackages: 99 | - m3 100 | - m3/customtransports 101 | - m3/thrift 102 | - m3/thriftudp 103 | - multi 104 | - prometheus 105 | - name: go.uber.org/atomic 106 | version: 1ea20fb1cbb1cc08cbd0d913a96dead89aa18289 107 | - name: go.uber.org/multierr 108 | version: 3c4937480c32f4c13a875a1829af76c98ca3d40a 109 | - name: go.uber.org/zap 110 | version: eeedf312bc6c57391d84767a4cd413f02a917974 111 | subpackages: 112 | - buffer 113 | - internal/bufferpool 114 | - internal/color 115 | - internal/exit 116 | - zapcore 117 | - name: golang.org/x/net 118 | version: 054b33e6527139ad5b1ec2f6232c3b175bd9a30c 119 | subpackages: 120 | - context 121 | - name: gopkg.in/yaml.v2 122 | version: 5420a8b6744d3b0345ab293f6fcba19c978f1183 123 | testImports: [] 124 | -------------------------------------------------------------------------------- /producer/ref_counted.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2018 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 producer 22 | 23 | import ( 24 | "sync" 25 | 26 | "go.uber.org/atomic" 27 | ) 28 | 29 | // OnFinalizeFn will be called when the message is being finalized. 30 | type OnFinalizeFn func(rm *RefCountedMessage) 31 | 32 | // RefCountedMessage is a reference counted message. 33 | type RefCountedMessage struct { 34 | sync.RWMutex 35 | Message 36 | 37 | size uint64 38 | onFinalizeFn OnFinalizeFn 39 | 40 | refCount *atomic.Int32 41 | isDroppedOrConsumed *atomic.Bool 42 | } 43 | 44 | // NewRefCountedMessage creates RefCountedMessage. 45 | func NewRefCountedMessage(m Message, fn OnFinalizeFn) *RefCountedMessage { 46 | return &RefCountedMessage{ 47 | Message: m, 48 | refCount: atomic.NewInt32(0), 49 | size: uint64(m.Size()), 50 | onFinalizeFn: fn, 51 | isDroppedOrConsumed: atomic.NewBool(false), 52 | } 53 | } 54 | 55 | // Accept returns true if the message can be accepted by the filter. 56 | func (rm *RefCountedMessage) Accept(fn FilterFunc) bool { 57 | return fn(rm.Message) 58 | } 59 | 60 | // IncRef increments the ref count. 61 | func (rm *RefCountedMessage) IncRef() { 62 | rm.refCount.Inc() 63 | } 64 | 65 | // DecRef decrements the ref count. If the reference count became zero after 66 | // the call, the message will be finalized as consumed. 67 | func (rm *RefCountedMessage) DecRef() { 68 | rc := rm.refCount.Dec() 69 | if rc == 0 { 70 | rm.finalize(Consumed) 71 | } 72 | if rc < 0 { 73 | panic("invalid ref count") 74 | } 75 | } 76 | 77 | // IncReads increments the reads count. 78 | func (rm *RefCountedMessage) IncReads() { 79 | rm.RLock() 80 | } 81 | 82 | // DecReads decrements the reads count. 83 | func (rm *RefCountedMessage) DecReads() { 84 | rm.RUnlock() 85 | } 86 | 87 | // Size returns the size of the message. 88 | func (rm *RefCountedMessage) Size() uint64 { 89 | return rm.size 90 | } 91 | 92 | // Drop drops the message without waiting for it to be consumed. 93 | func (rm *RefCountedMessage) Drop() bool { 94 | return rm.finalize(Dropped) 95 | } 96 | 97 | // IsDroppedOrConsumed returns true if the message has been dropped or consumed. 98 | func (rm *RefCountedMessage) IsDroppedOrConsumed() bool { 99 | return rm.isDroppedOrConsumed.Load() 100 | } 101 | 102 | func (rm *RefCountedMessage) finalize(r FinalizeReason) bool { 103 | // NB: This lock prevents the message from being finalized when its still 104 | // being read. 105 | rm.Lock() 106 | if rm.isDroppedOrConsumed.Load() { 107 | rm.Unlock() 108 | return false 109 | } 110 | rm.isDroppedOrConsumed.Store(true) 111 | rm.Unlock() 112 | if rm.onFinalizeFn != nil { 113 | rm.onFinalizeFn(rm) 114 | } 115 | rm.Message.Finalize(r) 116 | return true 117 | } 118 | -------------------------------------------------------------------------------- /producer/writer/options_test.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2018 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 writer 22 | 23 | import ( 24 | "testing" 25 | "time" 26 | 27 | "github.com/m3db/m3x/instrument" 28 | 29 | "github.com/stretchr/testify/require" 30 | ) 31 | 32 | func TestOptions(t *testing.T) { 33 | opts := NewOptions() 34 | 35 | require.Empty(t, opts.TopicName()) 36 | require.Equal(t, "topic", opts.SetTopicName("topic").TopicName()) 37 | 38 | require.Equal(t, defaultTopicWatchInitTimeout, opts.TopicWatchInitTimeout()) 39 | require.Equal(t, time.Second, opts.SetTopicWatchInitTimeout(time.Second).TopicWatchInitTimeout()) 40 | 41 | require.Equal(t, defaultPlacementWatchInitTimeout, opts.PlacementWatchInitTimeout()) 42 | require.Equal(t, time.Second, opts.SetPlacementWatchInitTimeout(time.Second).PlacementWatchInitTimeout()) 43 | 44 | require.Equal(t, defaultMessageQueueNewWritesScanInterval, opts.MessageQueueNewWritesScanInterval()) 45 | require.Equal(t, time.Second, opts.SetMessageQueueNewWritesScanInterval(time.Second).MessageQueueNewWritesScanInterval()) 46 | 47 | require.Equal(t, defaultMessageQueueFullScanInterval, opts.MessageQueueFullScanInterval()) 48 | require.Equal(t, time.Minute, opts.SetMessageQueueFullScanInterval(time.Minute).MessageQueueFullScanInterval()) 49 | 50 | require.Nil(t, opts.MessagePoolOptions()) 51 | 52 | require.Equal(t, defaultInitialAckMapSize, opts.InitialAckMapSize()) 53 | require.Equal(t, 123, opts.SetInitialAckMapSize(123).InitialAckMapSize()) 54 | 55 | require.Equal(t, defaultCloseCheckInterval, opts.CloseCheckInterval()) 56 | require.Equal(t, time.Second, opts.SetCloseCheckInterval(time.Second).CloseCheckInterval()) 57 | 58 | require.Equal(t, instrument.NewOptions(), opts.InstrumentOptions()) 59 | require.Nil(t, opts.SetInstrumentOptions(nil).InstrumentOptions()) 60 | } 61 | 62 | func TestConnectionOptions(t *testing.T) { 63 | opts := NewConnectionOptions() 64 | 65 | require.Equal(t, defaultConnectionDialTimeout, opts.DialTimeout()) 66 | require.Equal(t, time.Second, opts.SetDialTimeout(time.Second).DialTimeout()) 67 | 68 | require.Equal(t, defaultConnectionWriteTimeout, opts.WriteTimeout()) 69 | require.Equal(t, time.Second, opts.SetWriteTimeout(time.Second).WriteTimeout()) 70 | 71 | require.Equal(t, defaultConnectionResetDelay, opts.ResetDelay()) 72 | require.Equal(t, time.Second, opts.SetResetDelay(time.Second).ResetDelay()) 73 | 74 | require.Nil(t, opts.SetRetryOptions(nil).RetryOptions()) 75 | 76 | require.Equal(t, defaultConnectionFlushInterval, opts.FlushInterval()) 77 | require.Equal(t, 2*time.Second, opts.SetFlushInterval(2*time.Second).FlushInterval()) 78 | 79 | require.Equal(t, defaultConnectionBufferSize, opts.WriteBufferSize()) 80 | require.Equal(t, 1, opts.SetWriteBufferSize(1).WriteBufferSize()) 81 | 82 | require.Equal(t, defaultConnectionBufferSize, opts.ReadBufferSize()) 83 | require.Equal(t, 1, opts.SetReadBufferSize(1).ReadBufferSize()) 84 | 85 | } 86 | -------------------------------------------------------------------------------- /producer/writer/message.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2018 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 writer 22 | 23 | import ( 24 | "github.com/m3db/m3msg/generated/proto/msgpb" 25 | "github.com/m3db/m3msg/producer" 26 | "github.com/m3db/m3msg/protocol/proto" 27 | 28 | "go.uber.org/atomic" 29 | ) 30 | 31 | type message struct { 32 | *producer.RefCountedMessage 33 | 34 | pb msgpb.Message 35 | meta metadata 36 | initNanos int64 37 | retryAtNanos int64 38 | retried int 39 | // NB(cw) isAcked could be accessed concurrently by the background thread 40 | // in message writer and acked by consumer service writers. 41 | isAcked *atomic.Bool 42 | } 43 | 44 | func newMessage() *message { 45 | return &message{ 46 | retryAtNanos: 0, 47 | retried: 0, 48 | isAcked: atomic.NewBool(false), 49 | } 50 | } 51 | 52 | // Set sets the message. 53 | func (m *message) Set(meta metadata, rm *producer.RefCountedMessage, initNanos int64) { 54 | m.initNanos = initNanos 55 | m.meta = meta 56 | m.RefCountedMessage = rm 57 | m.ToProto(&m.pb) 58 | } 59 | 60 | // Close resets the states of the message. 61 | func (m *message) Close() { 62 | m.retryAtNanos = 0 63 | m.retried = 0 64 | m.isAcked.Store(false) 65 | m.ResetProto(&m.pb) 66 | } 67 | 68 | // InitNanos returns the nanosecond when the message was initiated. 69 | func (m *message) InitNanos() int64 { 70 | return m.initNanos 71 | } 72 | 73 | // RetryAtNanos returns the timestamp for next retry in nano seconds. 74 | func (m *message) RetryAtNanos() int64 { 75 | return m.retryAtNanos 76 | } 77 | 78 | // SetRetryAtNanos sets the next retry nanos. 79 | func (m *message) SetRetryAtNanos(value int64) { 80 | m.retryAtNanos = value 81 | } 82 | 83 | // WriteTimes returns the times the message has been written. 84 | func (m *message) WriteTimes() int { 85 | return m.retried 86 | } 87 | 88 | // IncWriteTimes increments the times the message has been written. 89 | func (m *message) IncWriteTimes() { 90 | m.retried++ 91 | } 92 | 93 | // IsAcked returns true if the message has been acked. 94 | func (m *message) IsAcked() bool { 95 | return m.isAcked.Load() 96 | } 97 | 98 | // Ack acknowledges the message. Duplicated acks on the same message might cause panic. 99 | func (m *message) Ack() { 100 | m.isAcked.Store(true) 101 | m.RefCountedMessage.DecRef() 102 | } 103 | 104 | // Metadata returns the metadata. 105 | func (m *message) Metadata() metadata { 106 | return m.meta 107 | } 108 | 109 | // Marshaler returns the marshaler and a bool to indicate whether the marshaler is valid. 110 | func (m *message) Marshaler() (proto.Marshaler, bool) { 111 | return &m.pb, !m.RefCountedMessage.IsDroppedOrConsumed() 112 | } 113 | 114 | func (m *message) ToProto(pb *msgpb.Message) { 115 | m.meta.ToProto(&pb.Metadata) 116 | pb.Value = m.RefCountedMessage.Bytes() 117 | } 118 | 119 | func (m *message) ResetProto(pb *msgpb.Message) { 120 | pb.Value = nil 121 | } 122 | -------------------------------------------------------------------------------- /topic/service_test.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2018 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 topic 22 | 23 | import ( 24 | "testing" 25 | 26 | "github.com/m3db/m3cluster/client" 27 | "github.com/m3db/m3cluster/kv" 28 | "github.com/m3db/m3cluster/kv/mem" 29 | "github.com/m3db/m3cluster/services" 30 | "github.com/m3db/m3msg/generated/proto/msgpb" 31 | 32 | "github.com/golang/mock/gomock" 33 | "github.com/stretchr/testify/require" 34 | ) 35 | 36 | func TestTopicService(t *testing.T) { 37 | ctrl := gomock.NewController(t) 38 | defer ctrl.Finish() 39 | 40 | kvOpts := kv.NewOverrideOptions().SetNamespace("foo") 41 | cs := client.NewMockClient(ctrl) 42 | store := mem.NewStore() 43 | cs.EXPECT().Store(kvOpts).Return(store, nil) 44 | 45 | s, err := NewService(NewServiceOptions().SetConfigService(cs).SetKVOverrideOptions(kvOpts)) 46 | require.NoError(t, err) 47 | 48 | topicName := "topic1" 49 | _, err = s.Get(topicName) 50 | require.Error(t, err) 51 | 52 | _, err = s.CheckAndSet(NewTopic(), kv.UninitializedVersion) 53 | require.Error(t, err) 54 | require.Contains(t, err.Error(), "invalid topic") 55 | 56 | w, err := s.Watch(topicName) 57 | require.NoError(t, err) 58 | require.Equal(t, 0, len(w.C())) 59 | 60 | topic1 := NewTopic(). 61 | SetName(topicName). 62 | SetNumberOfShards(100). 63 | SetConsumerServices([]ConsumerService{ 64 | NewConsumerService().SetConsumptionType(Shared).SetServiceID(services.NewServiceID().SetName("s1")), 65 | NewConsumerService().SetConsumptionType(Replicated).SetServiceID(services.NewServiceID().SetName("s2")), 66 | }) 67 | topic1, err = s.CheckAndSet(topic1, kv.UninitializedVersion) 68 | require.NoError(t, err) 69 | topic1, err = s.CheckAndSet(topic1, 1) 70 | require.NoError(t, err) 71 | 72 | topic2, err := s.Get(topicName) 73 | require.NoError(t, err) 74 | require.Equal(t, 2, topic2.Version()) 75 | require.Equal(t, topic1.Name(), topic2.Name()) 76 | require.Equal(t, topic1.NumberOfShards(), topic2.NumberOfShards()) 77 | require.Equal(t, topic1.ConsumerServices(), topic2.ConsumerServices()) 78 | 79 | <-w.C() 80 | topic, err := w.Get() 81 | require.NoError(t, err) 82 | require.Equal(t, topic2, topic) 83 | 84 | err = s.Delete(topicName) 85 | require.NoError(t, err) 86 | 87 | <-w.C() 88 | _, err = w.Get() 89 | require.Error(t, err) 90 | 91 | _, err = s.CheckAndSet(topic1, kv.UninitializedVersion) 92 | require.NoError(t, err) 93 | 94 | <-w.C() 95 | topic, err = w.Get() 96 | require.NoError(t, err) 97 | 98 | topic3, err := s.Get(topicName) 99 | require.NoError(t, err) 100 | require.Equal(t, 1, topic3.Version()) 101 | require.Equal(t, topic3, topic) 102 | 103 | version, err := store.Set(key(topicName), &msgpb.Message{Value: []byte("bad proto")}) 104 | require.NoError(t, err) 105 | require.Equal(t, 2, version) 106 | 107 | _, err = s.Get(topicName) 108 | require.Error(t, err) 109 | 110 | <-w.C() 111 | _, err = w.Get() 112 | require.Error(t, err) 113 | 114 | w.Close() 115 | } 116 | -------------------------------------------------------------------------------- /topic/topic_mock.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2018 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 | // Code generated by MockGen. DO NOT EDIT. 22 | // Source: github.com/m3db/m3msg/topic (interfaces: Service) 23 | 24 | // Package topic is a generated GoMock package. 25 | package topic 26 | 27 | import ( 28 | "reflect" 29 | 30 | "github.com/golang/mock/gomock" 31 | ) 32 | 33 | // MockService is a mock of Service interface 34 | type MockService struct { 35 | ctrl *gomock.Controller 36 | recorder *MockServiceMockRecorder 37 | } 38 | 39 | // MockServiceMockRecorder is the mock recorder for MockService 40 | type MockServiceMockRecorder struct { 41 | mock *MockService 42 | } 43 | 44 | // NewMockService creates a new mock instance 45 | func NewMockService(ctrl *gomock.Controller) *MockService { 46 | mock := &MockService{ctrl: ctrl} 47 | mock.recorder = &MockServiceMockRecorder{mock} 48 | return mock 49 | } 50 | 51 | // EXPECT returns an object that allows the caller to indicate expected use 52 | func (m *MockService) EXPECT() *MockServiceMockRecorder { 53 | return m.recorder 54 | } 55 | 56 | // CheckAndSet mocks base method 57 | func (m *MockService) CheckAndSet(arg0 Topic, arg1 int) (Topic, error) { 58 | ret := m.ctrl.Call(m, "CheckAndSet", arg0, arg1) 59 | ret0, _ := ret[0].(Topic) 60 | ret1, _ := ret[1].(error) 61 | return ret0, ret1 62 | } 63 | 64 | // CheckAndSet indicates an expected call of CheckAndSet 65 | func (mr *MockServiceMockRecorder) CheckAndSet(arg0, arg1 interface{}) *gomock.Call { 66 | return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CheckAndSet", reflect.TypeOf((*MockService)(nil).CheckAndSet), arg0, arg1) 67 | } 68 | 69 | // Delete mocks base method 70 | func (m *MockService) Delete(arg0 string) error { 71 | ret := m.ctrl.Call(m, "Delete", arg0) 72 | ret0, _ := ret[0].(error) 73 | return ret0 74 | } 75 | 76 | // Delete indicates an expected call of Delete 77 | func (mr *MockServiceMockRecorder) Delete(arg0 interface{}) *gomock.Call { 78 | return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Delete", reflect.TypeOf((*MockService)(nil).Delete), arg0) 79 | } 80 | 81 | // Get mocks base method 82 | func (m *MockService) Get(arg0 string) (Topic, error) { 83 | ret := m.ctrl.Call(m, "Get", arg0) 84 | ret0, _ := ret[0].(Topic) 85 | ret1, _ := ret[1].(error) 86 | return ret0, ret1 87 | } 88 | 89 | // Get indicates an expected call of Get 90 | func (mr *MockServiceMockRecorder) Get(arg0 interface{}) *gomock.Call { 91 | return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Get", reflect.TypeOf((*MockService)(nil).Get), arg0) 92 | } 93 | 94 | // Watch mocks base method 95 | func (m *MockService) Watch(arg0 string) (Watch, error) { 96 | ret := m.ctrl.Call(m, "Watch", arg0) 97 | ret0, _ := ret[0].(Watch) 98 | ret1, _ := ret[1].(error) 99 | return ret0, ret1 100 | } 101 | 102 | // Watch indicates an expected call of Watch 103 | func (mr *MockServiceMockRecorder) Watch(arg0 interface{}) *gomock.Call { 104 | return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Watch", reflect.TypeOf((*MockService)(nil).Watch), arg0) 105 | } 106 | -------------------------------------------------------------------------------- /producer/buffer/types.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2018 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 buffer 22 | 23 | import ( 24 | "time" 25 | 26 | "github.com/m3db/m3x/instrument" 27 | "github.com/m3db/m3x/retry" 28 | ) 29 | 30 | // OnFullStrategy defines the buffer behavior when the buffer is full. 31 | type OnFullStrategy string 32 | 33 | const ( 34 | // ReturnError means an error will be returned 35 | // on new buffer requests when the buffer is full. 36 | ReturnError OnFullStrategy = "returnError" 37 | 38 | // DropOldest means the oldest message in the buffer 39 | // will be dropped to make room for new buffer requests 40 | // when the buffer is full. 41 | DropOldest OnFullStrategy = "dropOldest" 42 | ) 43 | 44 | // Options configs the buffer. 45 | type Options interface { 46 | // OnFullStrategy returns the strategy when buffer is full. 47 | OnFullStrategy() OnFullStrategy 48 | 49 | // SetOnFullStrategy sets the strategy when buffer is full. 50 | SetOnFullStrategy(value OnFullStrategy) Options 51 | 52 | // MaxMessageSize returns the max message size. 53 | MaxMessageSize() int 54 | 55 | // SetMaxMessageSize sets the max message size. 56 | SetMaxMessageSize(value int) Options 57 | 58 | // MaxBufferSize returns the max buffer size. 59 | MaxBufferSize() int 60 | 61 | // SetMaxBufferSize sets the max buffer size. 62 | SetMaxBufferSize(value int) Options 63 | 64 | // CloseCheckInterval returns the close check interval. 65 | CloseCheckInterval() time.Duration 66 | 67 | // SetCloseCheckInterval sets the close check interval. 68 | SetCloseCheckInterval(value time.Duration) Options 69 | 70 | // DropOldestInterval returns the interval to drop oldest buffer. 71 | // The max buffer size might be spilled over during the interval. 72 | DropOldestInterval() time.Duration 73 | 74 | // SetDropOldestInterval sets the interval to drop oldest buffer. 75 | // The max buffer size might be spilled over during the interval. 76 | SetDropOldestInterval(value time.Duration) Options 77 | 78 | // ScanBatchSize returns the scan batch size. 79 | ScanBatchSize() int 80 | 81 | // SetScanBatchSize sets the scan batch size. 82 | SetScanBatchSize(value int) Options 83 | 84 | // AllowedSpilloverRatio returns the ratio for allowed buffer spill over, 85 | // below which the buffer will drop oldest messages asynchronizely for 86 | // better performance. When the limit for allowed spill over is reached, 87 | // the buffer will start to drop oldest messages synchronizely. 88 | AllowedSpilloverRatio() float64 89 | 90 | // SetAllowedSpilloverRatio sets the ratio for allowed buffer spill over. 91 | SetAllowedSpilloverRatio(value float64) Options 92 | 93 | // CleanupRetryOptions returns the cleanup retry options. 94 | CleanupRetryOptions() retry.Options 95 | 96 | // SetCleanupRetryOptions sets the cleanup retry options. 97 | SetCleanupRetryOptions(value retry.Options) Options 98 | 99 | // InstrumentOptions returns the instrument options. 100 | InstrumentOptions() instrument.Options 101 | 102 | // SetInstrumentOptions sets the instrument options. 103 | SetInstrumentOptions(value instrument.Options) Options 104 | 105 | // Validate validates the options. 106 | Validate() error 107 | } 108 | -------------------------------------------------------------------------------- /Makefile: -------------------------------------------------------------------------------- 1 | SELF_DIR := $(dir $(lastword $(MAKEFILE_LIST))) 2 | include $(SELF_DIR)/.ci/common.mk 3 | 4 | SHELL=/bin/bash -o pipefail 5 | 6 | html_report := coverage.html 7 | test := .ci/test-cover.sh 8 | convert-test-data := .ci/convert-test-data.sh 9 | coverfile := cover.out 10 | coverage_xml := coverage.xml 11 | junit_xml := junit.xml 12 | test_log := test.log 13 | lint_check := .ci/lint.sh 14 | metalint_check := .ci/metalint.sh 15 | metalint_config := .metalinter.json 16 | metalint_exclude := .excludemetalint 17 | package_root := github.com/m3db/m3msg 18 | gopath_prefix := $(GOPATH)/src 19 | vendor_prefix := vendor 20 | mockgen_package := github.com/golang/mock/mockgen 21 | mocks_output_dir := generated/mocks/mocks 22 | mocks_rules_dir := generated/mocks 23 | protoc_go_package := github.com/golang/protobuf/protoc-gen-go 24 | proto_output_dir := generated/proto 25 | proto_rules_dir := generated/proto 26 | auto_gen := .ci/auto-gen.sh 27 | license_dir := .ci/uber-licence 28 | license_node_modules := $(license_dir)/node_modules 29 | 30 | BUILD := $(abspath ./bin) 31 | LINUX_AMD64_ENV := GOOS=linux GOARCH=amd64 CGO_ENABLED=0 32 | VENDOR_ENV := GO15VENDOREXPERIMENT=1 33 | 34 | .PHONY: setup 35 | setup: 36 | mkdir -p $(BUILD) 37 | 38 | .PHONY: lint 39 | lint: 40 | @which golint > /dev/null || go get -u github.com/golang/lint/golint 41 | $(VENDOR_ENV) $(lint_check) 42 | 43 | .PHONY: metalint 44 | metalint: install-metalinter install-linter-badtime 45 | @($(metalint_check) $(metalint_config) $(metalint_exclude) && echo "metalinted successfully!") || (echo "metalinter failed" && exit 1) 46 | 47 | .PHONY: test-internal 48 | test-internal: 49 | @which go-junit-report > /dev/null || go get -u github.com/sectioneight/go-junit-report 50 | @$(VENDOR_ENV) $(test) $(coverfile) | tee $(test_log) 51 | 52 | .PHONY: test-integration 53 | test-integration: 54 | go test -v -tags=integration ./integration 55 | 56 | .PHONY: test-xml 57 | test-xml: test-internal 58 | go-junit-report < $(test_log) > $(junit_xml) 59 | gocov convert $(coverfile) | gocov-xml > $(coverage_xml) 60 | @$(convert-test-data) $(coverage_xml) 61 | @rm $(coverfile) &> /dev/null 62 | 63 | .PHONY: test 64 | test: test-internal 65 | gocov convert $(coverfile) | gocov report 66 | 67 | .PHONY: testhtml 68 | testhtml: test-internal 69 | gocov convert $(coverfile) | gocov-html > $(html_report) && open $(html_report) 70 | @rm -f $(test_log) &> /dev/null 71 | 72 | .PHONY: test-ci-unit 73 | test-ci-unit: test-internal 74 | @which goveralls > /dev/null || go get -u -f github.com/mattn/goveralls 75 | goveralls -coverprofile=$(coverfile) -service=travis-ci || echo -e "\x1b[31mCoveralls failed\x1b[m" 76 | 77 | .PHONY: test-ci-integration 78 | test-ci-integration: 79 | $(test_ci_integration) 80 | 81 | .PHONY: install-mockgen 82 | install-mockgen: install-vendor 83 | @echo Installing mockgen 84 | glide install 85 | 86 | .PHONY: install-licence-bin 87 | install-license-bin: install-vendor 88 | @echo Installing node modules 89 | [ -d $(license_node_modules) ] || (cd $(license_dir) && npm install) 90 | 91 | .PHONY: install-proto-bin 92 | install-proto-bin: install-vendor 93 | @echo Installing protobuf binaries 94 | @echo Note: the protobuf compiler v3.0.0 can be downloaded from https://github.com/google/protobuf/releases or built from source at https://github.com/google/protobuf. 95 | go install $(package_root)/$(vendor_prefix)/$(protoc_go_package) 96 | 97 | .PHONY: mock-gen 98 | mock-gen: install-mockgen install-license-bin install-util-mockclean 99 | @echo Generating mocks 100 | PACKAGE=$(package_root) $(auto_gen) $(mocks_output_dir) $(mocks_rules_dir) 101 | 102 | .PHONY: mock-gen-deps 103 | mock-gen-no-deps: 104 | @echo Generating mocks 105 | PACKAGE=$(package_root) $(auto_gen) $(mocks_output_dir) $(mocks_rules_dir) 106 | 107 | .PHONY: proto-gen 108 | proto-gen: install-proto-bin install-license-bin 109 | @echo Generating protobuf files 110 | PACKAGE=$(package_root) $(auto_gen) $(proto_output_dir) $(proto_rules_dir) 111 | 112 | .PHONY: all-gen 113 | all-gen: 114 | 115 | .PHONY: clean 116 | clean: 117 | @rm -f *.html *.xml *.out *.test 118 | 119 | .PHONY: all 120 | all: metalint test-ci-unit test-ci-integration 121 | @echo make all successfully finished 122 | 123 | .DEFAULT_GOAL := all 124 | -------------------------------------------------------------------------------- /producer/writer/shard_writer_mock.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2018 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 | // Code generated by MockGen. DO NOT EDIT. 22 | // Source: github.com/m3db/m3msg/producer/writer/shard_go 23 | 24 | // Package writer is a generated GoMock package. 25 | package writer 26 | 27 | import ( 28 | "reflect" 29 | 30 | "github.com/m3db/m3cluster/placement" 31 | "github.com/m3db/m3msg/producer" 32 | 33 | "github.com/golang/mock/gomock" 34 | ) 35 | 36 | // MockshardWriter is a mock of shardWriter interface 37 | type MockshardWriter struct { 38 | ctrl *gomock.Controller 39 | recorder *MockshardWriterMockRecorder 40 | } 41 | 42 | // MockshardWriterMockRecorder is the mock recorder for MockshardWriter 43 | type MockshardWriterMockRecorder struct { 44 | mock *MockshardWriter 45 | } 46 | 47 | // NewMockshardWriter creates a new mock instance 48 | func NewMockshardWriter(ctrl *gomock.Controller) *MockshardWriter { 49 | mock := &MockshardWriter{ctrl: ctrl} 50 | mock.recorder = &MockshardWriterMockRecorder{mock} 51 | return mock 52 | } 53 | 54 | // EXPECT returns an object that allows the caller to indicate expected use 55 | func (m *MockshardWriter) EXPECT() *MockshardWriterMockRecorder { 56 | return m.recorder 57 | } 58 | 59 | // Write mocks base method 60 | func (m *MockshardWriter) Write(rm *producer.RefCountedMessage) { 61 | m.ctrl.Call(m, "Write", rm) 62 | } 63 | 64 | // Write indicates an expected call of Write 65 | func (mr *MockshardWriterMockRecorder) Write(rm interface{}) *gomock.Call { 66 | return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Write", reflect.TypeOf((*MockshardWriter)(nil).Write), rm) 67 | } 68 | 69 | // UpdateInstances mocks base method 70 | func (m *MockshardWriter) UpdateInstances(instances []placement.Instance, cws map[string]consumerWriter) { 71 | m.ctrl.Call(m, "UpdateInstances", instances, cws) 72 | } 73 | 74 | // UpdateInstances indicates an expected call of UpdateInstances 75 | func (mr *MockshardWriterMockRecorder) UpdateInstances(instances, cws interface{}) *gomock.Call { 76 | return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateInstances", reflect.TypeOf((*MockshardWriter)(nil).UpdateInstances), instances, cws) 77 | } 78 | 79 | // SetMessageTTLNanos mocks base method 80 | func (m *MockshardWriter) SetMessageTTLNanos(value int64) { 81 | m.ctrl.Call(m, "SetMessageTTLNanos", value) 82 | } 83 | 84 | // SetMessageTTLNanos indicates an expected call of SetMessageTTLNanos 85 | func (mr *MockshardWriterMockRecorder) SetMessageTTLNanos(value interface{}) *gomock.Call { 86 | return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SetMessageTTLNanos", reflect.TypeOf((*MockshardWriter)(nil).SetMessageTTLNanos), value) 87 | } 88 | 89 | // Close mocks base method 90 | func (m *MockshardWriter) Close() { 91 | m.ctrl.Call(m, "Close") 92 | } 93 | 94 | // Close indicates an expected call of Close 95 | func (mr *MockshardWriterMockRecorder) Close() *gomock.Call { 96 | return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Close", reflect.TypeOf((*MockshardWriter)(nil).Close)) 97 | } 98 | 99 | // QueueSize mocks base method 100 | func (m *MockshardWriter) QueueSize() int { 101 | ret := m.ctrl.Call(m, "QueueSize") 102 | ret0, _ := ret[0].(int) 103 | return ret0 104 | } 105 | 106 | // QueueSize indicates an expected call of QueueSize 107 | func (mr *MockshardWriterMockRecorder) QueueSize() *gomock.Call { 108 | return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "QueueSize", reflect.TypeOf((*MockshardWriter)(nil).QueueSize)) 109 | } 110 | -------------------------------------------------------------------------------- /protocol/proto/proto_mock.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2018 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 | // Code generated by MockGen. DO NOT EDIT. 22 | // Source: github.com/m3db/m3msg/protocol/proto (interfaces: Encoder,Decoder) 23 | 24 | // Package proto is a generated GoMock package. 25 | package proto 26 | 27 | import ( 28 | "io" 29 | "reflect" 30 | 31 | "github.com/golang/mock/gomock" 32 | ) 33 | 34 | // MockEncoder is a mock of Encoder interface 35 | type MockEncoder struct { 36 | ctrl *gomock.Controller 37 | recorder *MockEncoderMockRecorder 38 | } 39 | 40 | // MockEncoderMockRecorder is the mock recorder for MockEncoder 41 | type MockEncoderMockRecorder struct { 42 | mock *MockEncoder 43 | } 44 | 45 | // NewMockEncoder creates a new mock instance 46 | func NewMockEncoder(ctrl *gomock.Controller) *MockEncoder { 47 | mock := &MockEncoder{ctrl: ctrl} 48 | mock.recorder = &MockEncoderMockRecorder{mock} 49 | return mock 50 | } 51 | 52 | // EXPECT returns an object that allows the caller to indicate expected use 53 | func (m *MockEncoder) EXPECT() *MockEncoderMockRecorder { 54 | return m.recorder 55 | } 56 | 57 | // Bytes mocks base method 58 | func (m *MockEncoder) Bytes() []byte { 59 | ret := m.ctrl.Call(m, "Bytes") 60 | ret0, _ := ret[0].([]byte) 61 | return ret0 62 | } 63 | 64 | // Bytes indicates an expected call of Bytes 65 | func (mr *MockEncoderMockRecorder) Bytes() *gomock.Call { 66 | return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Bytes", reflect.TypeOf((*MockEncoder)(nil).Bytes)) 67 | } 68 | 69 | // Encode mocks base method 70 | func (m *MockEncoder) Encode(arg0 Marshaler) error { 71 | ret := m.ctrl.Call(m, "Encode", arg0) 72 | ret0, _ := ret[0].(error) 73 | return ret0 74 | } 75 | 76 | // Encode indicates an expected call of Encode 77 | func (mr *MockEncoderMockRecorder) Encode(arg0 interface{}) *gomock.Call { 78 | return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Encode", reflect.TypeOf((*MockEncoder)(nil).Encode), arg0) 79 | } 80 | 81 | // MockDecoder is a mock of Decoder interface 82 | type MockDecoder struct { 83 | ctrl *gomock.Controller 84 | recorder *MockDecoderMockRecorder 85 | } 86 | 87 | // MockDecoderMockRecorder is the mock recorder for MockDecoder 88 | type MockDecoderMockRecorder struct { 89 | mock *MockDecoder 90 | } 91 | 92 | // NewMockDecoder creates a new mock instance 93 | func NewMockDecoder(ctrl *gomock.Controller) *MockDecoder { 94 | mock := &MockDecoder{ctrl: ctrl} 95 | mock.recorder = &MockDecoderMockRecorder{mock} 96 | return mock 97 | } 98 | 99 | // EXPECT returns an object that allows the caller to indicate expected use 100 | func (m *MockDecoder) EXPECT() *MockDecoderMockRecorder { 101 | return m.recorder 102 | } 103 | 104 | // Decode mocks base method 105 | func (m *MockDecoder) Decode(arg0 Unmarshaler) error { 106 | ret := m.ctrl.Call(m, "Decode", arg0) 107 | ret0, _ := ret[0].(error) 108 | return ret0 109 | } 110 | 111 | // Decode indicates an expected call of Decode 112 | func (mr *MockDecoderMockRecorder) Decode(arg0 interface{}) *gomock.Call { 113 | return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Decode", reflect.TypeOf((*MockDecoder)(nil).Decode), arg0) 114 | } 115 | 116 | // ResetReader mocks base method 117 | func (m *MockDecoder) ResetReader(arg0 io.Reader) { 118 | m.ctrl.Call(m, "ResetReader", arg0) 119 | } 120 | 121 | // ResetReader indicates an expected call of ResetReader 122 | func (mr *MockDecoderMockRecorder) ResetReader(arg0 interface{}) *gomock.Call { 123 | return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ResetReader", reflect.TypeOf((*MockDecoder)(nil).ResetReader), arg0) 124 | } 125 | -------------------------------------------------------------------------------- /producer/config/writer_test.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2018 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 config 22 | 23 | import ( 24 | "testing" 25 | "time" 26 | 27 | "github.com/m3db/m3cluster/client" 28 | "github.com/m3db/m3cluster/kv" 29 | "github.com/m3db/m3cluster/services" 30 | "github.com/m3db/m3x/instrument" 31 | 32 | "github.com/golang/mock/gomock" 33 | "github.com/stretchr/testify/require" 34 | yaml "gopkg.in/yaml.v2" 35 | ) 36 | 37 | func TestConnectionConfiguration(t *testing.T) { 38 | str := ` 39 | dialTimeout: 3s 40 | writeTimeout: 2s 41 | keepAlivePeriod: 20s 42 | resetDelay: 1s 43 | retry: 44 | initialBackoff: 1ms 45 | maxBackoff: 2ms 46 | flushInterval: 2s 47 | writeBufferSize: 100 48 | readBufferSize: 200 49 | ` 50 | 51 | var cfg ConnectionConfiguration 52 | require.NoError(t, yaml.Unmarshal([]byte(str), &cfg)) 53 | 54 | cOpts := cfg.NewOptions(instrument.NewOptions()) 55 | require.Equal(t, 3*time.Second, cOpts.DialTimeout()) 56 | require.Equal(t, 2*time.Second, cOpts.WriteTimeout()) 57 | require.Equal(t, 20*time.Second, cOpts.KeepAlivePeriod()) 58 | require.Equal(t, time.Second, cOpts.ResetDelay()) 59 | require.Equal(t, time.Millisecond, cOpts.RetryOptions().InitialBackoff()) 60 | require.Equal(t, 2*time.Millisecond, cOpts.RetryOptions().MaxBackoff()) 61 | require.Equal(t, 2*time.Second, cOpts.FlushInterval()) 62 | require.Equal(t, 100, cOpts.WriteBufferSize()) 63 | require.Equal(t, 200, cOpts.ReadBufferSize()) 64 | } 65 | 66 | func TestWriterConfiguration(t *testing.T) { 67 | str := ` 68 | topicName: testTopic 69 | topicServiceOverride: 70 | zone: z1 71 | namespace: n1 72 | topicWatchInitTimeout: 1s 73 | placementServiceOverride: 74 | namespaces: 75 | placement: n2 76 | placementWatchInitTimeout: 2s 77 | messagePool: 78 | size: 5 79 | messageRetry: 80 | initialBackoff: 1ms 81 | messageQueueNewWritesScanInterval: 200ms 82 | messageQueueFullScanInterval: 10s 83 | messageQueueScanBatchSize: 1024 84 | initialAckMapSize: 1024 85 | closeCheckInterval: 2s 86 | ackErrorRetry: 87 | initialBackoff: 2ms 88 | connection: 89 | dialTimeout: 5s 90 | encoder: 91 | maxMessageSize: 100 92 | decoder: 93 | maxMessageSize: 200 94 | ` 95 | var cfg WriterConfiguration 96 | require.NoError(t, yaml.Unmarshal([]byte(str), &cfg)) 97 | 98 | ctrl := gomock.NewController(t) 99 | defer ctrl.Finish() 100 | 101 | cs := client.NewMockClient(ctrl) 102 | cs.EXPECT().Store(kv.NewOverrideOptions().SetZone("z1").SetNamespace("n1")).Return(nil, nil) 103 | cs.EXPECT().Services( 104 | services.NewOverrideOptions().SetNamespaceOptions( 105 | services.NewNamespaceOptions().SetPlacementNamespace("n2"), 106 | ), 107 | ).Return(nil, nil) 108 | 109 | wOpts, err := cfg.NewOptions(cs, instrument.NewOptions()) 110 | require.NoError(t, err) 111 | require.Equal(t, "testTopic", wOpts.TopicName()) 112 | require.Equal(t, time.Second, wOpts.TopicWatchInitTimeout()) 113 | require.Equal(t, 2*time.Second, wOpts.PlacementWatchInitTimeout()) 114 | require.Equal(t, 5, wOpts.MessagePoolOptions().Size()) 115 | require.Equal(t, time.Millisecond, wOpts.MessageRetryOptions().InitialBackoff()) 116 | require.Equal(t, 200*time.Millisecond, wOpts.MessageQueueNewWritesScanInterval()) 117 | require.Equal(t, 10*time.Second, wOpts.MessageQueueFullScanInterval()) 118 | require.Equal(t, 1024, wOpts.MessageQueueScanBatchSize()) 119 | require.Equal(t, 1024, wOpts.InitialAckMapSize()) 120 | require.Equal(t, 2*time.Second, wOpts.CloseCheckInterval()) 121 | require.Equal(t, 2*time.Millisecond, wOpts.AckErrorRetryOptions().InitialBackoff()) 122 | require.Equal(t, 5*time.Second, wOpts.ConnectionOptions().DialTimeout()) 123 | require.Equal(t, 100, wOpts.EncoderOptions().MaxMessageSize()) 124 | require.Equal(t, 200, wOpts.DecoderOptions().MaxMessageSize()) 125 | } 126 | -------------------------------------------------------------------------------- /producer/types.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2018 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 producer 22 | 23 | import ( 24 | "github.com/m3db/m3cluster/services" 25 | ) 26 | 27 | // FinalizeReason defines the reason why the message is being finalized by Producer. 28 | type FinalizeReason int 29 | 30 | const ( 31 | // Consumed means the message has been fully consumed. 32 | Consumed FinalizeReason = iota 33 | 34 | // Dropped means the message has been dropped. 35 | Dropped 36 | ) 37 | 38 | // Message contains the data that will be produced by the producer. 39 | // It should only be finalized by the producer. 40 | type Message interface { 41 | // Shard returns the shard of the message. 42 | Shard() uint32 43 | 44 | // Bytes returns the bytes of the message. 45 | Bytes() []byte 46 | 47 | // Size returns the size of the bytes of the message. 48 | Size() int 49 | 50 | // Finalize will be called by producer to indicate the end of its lifecycle. 51 | Finalize(FinalizeReason) 52 | } 53 | 54 | // CloseType decides how the producer should be closed. 55 | type CloseType int 56 | 57 | const ( 58 | // WaitForConsumption blocks the close call until all the messages have been consumed. 59 | WaitForConsumption CloseType = iota 60 | // DropEverything will close the producer and drop all the messages that have not been consumed. 61 | DropEverything 62 | ) 63 | 64 | // Producer produces message to a topic. 65 | type Producer interface { 66 | // Produce produces the message. 67 | Produce(m Message) error 68 | 69 | // RegisterFilter registers a filter to a consumer service. 70 | RegisterFilter(sid services.ServiceID, fn FilterFunc) 71 | 72 | // UnregisterFilter unregisters the filter of a consumer service. 73 | UnregisterFilter(sid services.ServiceID) 74 | 75 | // Init initializes a producer. 76 | Init() error 77 | 78 | // Close stops the producer from accepting new requests immediately. 79 | // If the CloseType is WaitForConsumption, then it will block until all the messages have been consumed. 80 | // If the CloseType is DropEverything, then it will simply drop all the messages buffered and return. 81 | Close(ct CloseType) 82 | } 83 | 84 | // FilterFunc can filter message. 85 | type FilterFunc func(m Message) bool 86 | 87 | // Options configs a producer. 88 | type Options interface { 89 | // Buffer returns the buffer. 90 | Buffer() Buffer 91 | 92 | // SetBuffer sets the buffer. 93 | SetBuffer(value Buffer) Options 94 | 95 | // Writer returns the writer. 96 | Writer() Writer 97 | 98 | // SetWriter sets the writer. 99 | SetWriter(value Writer) Options 100 | } 101 | 102 | // Buffer buffers all the messages in the producer. 103 | type Buffer interface { 104 | // Add adds message to the buffer and returns a reference counted message. 105 | Add(m Message) (*RefCountedMessage, error) 106 | 107 | // Init initializes the buffer. 108 | Init() 109 | 110 | // Close stops the buffer from accepting new requests immediately. 111 | // If the CloseType is WaitForConsumption, then it will block until all the messages have been consumed. 112 | // If the CloseType is DropEverything, then it will simply drop all the messages buffered and return. 113 | Close(ct CloseType) 114 | } 115 | 116 | // Writer writes all the messages out to the consumer services. 117 | type Writer interface { 118 | // Write writes a reference counted message out. 119 | Write(rm *RefCountedMessage) error 120 | 121 | // RegisterFilter registers a filter to a consumer service. 122 | RegisterFilter(sid services.ServiceID, fn FilterFunc) 123 | 124 | // UnregisterFilter unregisters the filter of a consumer service. 125 | UnregisterFilter(sid services.ServiceID) 126 | 127 | // Init initializes a writer. 128 | Init() error 129 | 130 | // Close closes the writer. 131 | Close() 132 | } 133 | -------------------------------------------------------------------------------- /consumer/types.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2018 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 consumer 22 | 23 | import ( 24 | "net" 25 | "time" 26 | 27 | "github.com/m3db/m3msg/protocol/proto" 28 | "github.com/m3db/m3x/instrument" 29 | "github.com/m3db/m3x/pool" 30 | "github.com/m3db/m3x/server" 31 | ) 32 | 33 | // Message carries the data that needs to be processed. 34 | type Message interface { 35 | // Bytes returns the bytes. 36 | Bytes() []byte 37 | 38 | // Ack acks the message. 39 | Ack() 40 | } 41 | 42 | // Consumer receives messages from a connection. 43 | type Consumer interface { 44 | // Message waits for and returns the next message received. 45 | Message() (Message, error) 46 | 47 | // Init initializes the consumer. 48 | Init() 49 | 50 | // Close closes the consumer. 51 | Close() 52 | } 53 | 54 | // Listener is a consumer listener based on a network address. 55 | type Listener interface { 56 | // Accept waits for and returns the next connection based consumer. 57 | Accept() (Consumer, error) 58 | 59 | // Close closes the listener. 60 | // Any blocked Accept operations will be unblocked and return errors. 61 | Close() error 62 | 63 | // Addr returns the listener's network address. 64 | Addr() net.Addr 65 | } 66 | 67 | // Options configs the consumer listener. 68 | type Options interface { 69 | // EncoderOptions returns the options for Encoder. 70 | EncoderOptions() proto.Options 71 | 72 | // SetEncoderOptions sets the options for Encoder. 73 | SetEncoderOptions(value proto.Options) Options 74 | 75 | // DecoderOptions returns the options for Decoder. 76 | DecoderOptions() proto.Options 77 | 78 | // SetDecoderOptions sets the options for Decoder. 79 | SetDecoderOptions(value proto.Options) Options 80 | 81 | // MessagePoolOptions returns the options for message pool. 82 | MessagePoolOptions() pool.ObjectPoolOptions 83 | 84 | // SetMessagePoolOptions sets the options for message pool. 85 | SetMessagePoolOptions(value pool.ObjectPoolOptions) Options 86 | 87 | // AckFlushInterval returns the ack flush interval. 88 | AckFlushInterval() time.Duration 89 | 90 | // SetAckFlushInterval sets the ack flush interval. 91 | SetAckFlushInterval(value time.Duration) Options 92 | 93 | // AckBufferSize returns the ack buffer size. 94 | AckBufferSize() int 95 | 96 | // SetAckBufferSize sets the ack buffer size. 97 | SetAckBufferSize(value int) Options 98 | 99 | // ConnectionWriteBufferSize returns the size of buffer before a write or a read. 100 | ConnectionWriteBufferSize() int 101 | 102 | // SetConnectionWriteBufferSize sets the buffer size. 103 | SetConnectionWriteBufferSize(value int) Options 104 | 105 | // ConnectionReadBufferSize returns the size of buffer before a write or a read. 106 | ConnectionReadBufferSize() int 107 | 108 | // SetConnectionWriteBufferSize sets the buffer size. 109 | SetConnectionReadBufferSize(value int) Options 110 | 111 | // InstrumentOptions returns the instrument options. 112 | InstrumentOptions() instrument.Options 113 | 114 | // SetInstrumentOptions sets the instrument options. 115 | SetInstrumentOptions(value instrument.Options) Options 116 | } 117 | 118 | // ConsumeFn processes the consumer. 119 | type ConsumeFn func(c Consumer) 120 | 121 | // ServerOptions configs the consumer server. 122 | type ServerOptions interface { 123 | // ConsumeFn returns the ConsumeFn. 124 | ConsumeFn() ConsumeFn 125 | 126 | // SetConsumeFn sets the ConsumeFn. 127 | SetConsumeFn(value ConsumeFn) ServerOptions 128 | 129 | // RetryOptions returns the options for connection retrier. 130 | ServerOptions() server.Options 131 | 132 | // SetRetryOptions sets the options for connection retrier. 133 | SetServerOptions(value server.Options) ServerOptions 134 | 135 | // InstrumentOptions returns the instrument options. 136 | ConsumerOptions() Options 137 | 138 | // SetInstrumentOptions sets the instrument options. 139 | SetConsumerOptions(value Options) ServerOptions 140 | } 141 | -------------------------------------------------------------------------------- /producer/writer/consumer_service_writer_mock.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2018 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 | // Code generated by MockGen. DO NOT EDIT. 22 | // Source: github.com/m3db/m3msg/producer/writer/consumer_service_go 23 | 24 | // Package writer is a generated GoMock package. 25 | package writer 26 | 27 | import ( 28 | "reflect" 29 | 30 | "github.com/m3db/m3msg/producer" 31 | 32 | "github.com/golang/mock/gomock" 33 | ) 34 | 35 | // MockconsumerServiceWriter is a mock of consumerServiceWriter interface 36 | type MockconsumerServiceWriter struct { 37 | ctrl *gomock.Controller 38 | recorder *MockconsumerServiceWriterMockRecorder 39 | } 40 | 41 | // MockconsumerServiceWriterMockRecorder is the mock recorder for MockconsumerServiceWriter 42 | type MockconsumerServiceWriterMockRecorder struct { 43 | mock *MockconsumerServiceWriter 44 | } 45 | 46 | // NewMockconsumerServiceWriter creates a new mock instance 47 | func NewMockconsumerServiceWriter(ctrl *gomock.Controller) *MockconsumerServiceWriter { 48 | mock := &MockconsumerServiceWriter{ctrl: ctrl} 49 | mock.recorder = &MockconsumerServiceWriterMockRecorder{mock} 50 | return mock 51 | } 52 | 53 | // EXPECT returns an object that allows the caller to indicate expected use 54 | func (m *MockconsumerServiceWriter) EXPECT() *MockconsumerServiceWriterMockRecorder { 55 | return m.recorder 56 | } 57 | 58 | // Write mocks base method 59 | func (m *MockconsumerServiceWriter) Write(rm *producer.RefCountedMessage) { 60 | m.ctrl.Call(m, "Write", rm) 61 | } 62 | 63 | // Write indicates an expected call of Write 64 | func (mr *MockconsumerServiceWriterMockRecorder) Write(rm interface{}) *gomock.Call { 65 | return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Write", reflect.TypeOf((*MockconsumerServiceWriter)(nil).Write), rm) 66 | } 67 | 68 | // Init mocks base method 69 | func (m *MockconsumerServiceWriter) Init(arg0 initType) error { 70 | ret := m.ctrl.Call(m, "Init", arg0) 71 | ret0, _ := ret[0].(error) 72 | return ret0 73 | } 74 | 75 | // Init indicates an expected call of Init 76 | func (mr *MockconsumerServiceWriterMockRecorder) Init(arg0 interface{}) *gomock.Call { 77 | return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Init", reflect.TypeOf((*MockconsumerServiceWriter)(nil).Init), arg0) 78 | } 79 | 80 | // Close mocks base method 81 | func (m *MockconsumerServiceWriter) Close() { 82 | m.ctrl.Call(m, "Close") 83 | } 84 | 85 | // Close indicates an expected call of Close 86 | func (mr *MockconsumerServiceWriterMockRecorder) Close() *gomock.Call { 87 | return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Close", reflect.TypeOf((*MockconsumerServiceWriter)(nil).Close)) 88 | } 89 | 90 | // SetMessageTTLNanos mocks base method 91 | func (m *MockconsumerServiceWriter) SetMessageTTLNanos(value int64) { 92 | m.ctrl.Call(m, "SetMessageTTLNanos", value) 93 | } 94 | 95 | // SetMessageTTLNanos indicates an expected call of SetMessageTTLNanos 96 | func (mr *MockconsumerServiceWriterMockRecorder) SetMessageTTLNanos(value interface{}) *gomock.Call { 97 | return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SetMessageTTLNanos", reflect.TypeOf((*MockconsumerServiceWriter)(nil).SetMessageTTLNanos), value) 98 | } 99 | 100 | // RegisterFilter mocks base method 101 | func (m *MockconsumerServiceWriter) RegisterFilter(fn producer.FilterFunc) { 102 | m.ctrl.Call(m, "RegisterFilter", fn) 103 | } 104 | 105 | // RegisterFilter indicates an expected call of RegisterFilter 106 | func (mr *MockconsumerServiceWriterMockRecorder) RegisterFilter(fn interface{}) *gomock.Call { 107 | return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "RegisterFilter", reflect.TypeOf((*MockconsumerServiceWriter)(nil).RegisterFilter), fn) 108 | } 109 | 110 | // UnregisterFilter mocks base method 111 | func (m *MockconsumerServiceWriter) UnregisterFilter() { 112 | m.ctrl.Call(m, "UnregisterFilter") 113 | } 114 | 115 | // UnregisterFilter indicates an expected call of UnregisterFilter 116 | func (mr *MockconsumerServiceWriterMockRecorder) UnregisterFilter() *gomock.Call { 117 | return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UnregisterFilter", reflect.TypeOf((*MockconsumerServiceWriter)(nil).UnregisterFilter)) 118 | } 119 | -------------------------------------------------------------------------------- /producer/ref_counted_test.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2018 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 producer 22 | 23 | import ( 24 | "sync" 25 | "testing" 26 | "time" 27 | 28 | "github.com/golang/mock/gomock" 29 | "github.com/stretchr/testify/require" 30 | ) 31 | 32 | func TestRefCountedMessageConsume(t *testing.T) { 33 | ctrl := gomock.NewController(t) 34 | defer ctrl.Finish() 35 | 36 | mm := NewMockMessage(ctrl) 37 | mm.EXPECT().Size().Return(100).AnyTimes() 38 | mm.EXPECT().Finalize(Consumed) 39 | 40 | rm := NewRefCountedMessage(mm, nil) 41 | require.Equal(t, mm.Size(), int(rm.Size())) 42 | require.False(t, rm.IsDroppedOrConsumed()) 43 | 44 | rm.IncRef() 45 | rm.DecRef() 46 | require.True(t, rm.IsDroppedOrConsumed()) 47 | 48 | rm.IncRef() 49 | rm.DecRef() 50 | require.True(t, rm.IsDroppedOrConsumed()) 51 | 52 | rm.Drop() 53 | require.True(t, rm.IsDroppedOrConsumed()) 54 | } 55 | 56 | func TestRefCountedMessageDrop(t *testing.T) { 57 | ctrl := gomock.NewController(t) 58 | defer ctrl.Finish() 59 | 60 | mm := NewMockMessage(ctrl) 61 | mm.EXPECT().Size().Return(100).AnyTimes() 62 | mm.EXPECT().Finalize(Dropped) 63 | 64 | rm := NewRefCountedMessage(mm, nil) 65 | require.Equal(t, mm.Size(), int(rm.Size())) 66 | require.False(t, rm.IsDroppedOrConsumed()) 67 | 68 | rm.Drop() 69 | require.True(t, rm.IsDroppedOrConsumed()) 70 | 71 | rm.IncRef() 72 | rm.DecRef() 73 | require.True(t, rm.IsDroppedOrConsumed()) 74 | 75 | rm.Drop() 76 | require.True(t, rm.IsDroppedOrConsumed()) 77 | } 78 | 79 | func TestRefCountedMessageBytesReadBlocking(t *testing.T) { 80 | ctrl := gomock.NewController(t) 81 | defer ctrl.Finish() 82 | 83 | mm := NewMockMessage(ctrl) 84 | mockBytes := []byte("foo") 85 | mm.EXPECT().Size().Return(3) 86 | mm.EXPECT().Bytes().Return(mockBytes) 87 | 88 | rm := NewRefCountedMessage(mm, nil) 89 | rm.IncReads() 90 | b := rm.Bytes() 91 | require.Equal(t, mockBytes, b) 92 | require.False(t, rm.IsDroppedOrConsumed()) 93 | 94 | doneCh := make(chan struct{}) 95 | go func() { 96 | mm.EXPECT().Finalize(Dropped) 97 | rm.Drop() 98 | close(doneCh) 99 | }() 100 | 101 | select { 102 | case <-doneCh: 103 | require.FailNow(t, "not expected") 104 | case <-time.After(time.Second): 105 | } 106 | rm.DecReads() 107 | <-doneCh 108 | require.True(t, rm.IsDroppedOrConsumed()) 109 | } 110 | 111 | func TestRefCountedMessageDecPanic(t *testing.T) { 112 | ctrl := gomock.NewController(t) 113 | defer ctrl.Finish() 114 | 115 | mm := NewMockMessage(ctrl) 116 | mm.EXPECT().Size().Return(0) 117 | rm := NewRefCountedMessage(mm, nil) 118 | require.Panics(t, rm.DecRef) 119 | } 120 | 121 | func TestRefCountedMessageFilter(t *testing.T) { 122 | ctrl := gomock.NewController(t) 123 | defer ctrl.Finish() 124 | 125 | var called int 126 | filter := func(m Message) bool { 127 | called++ 128 | return m.Shard() == 0 129 | } 130 | 131 | mm := NewMockMessage(ctrl) 132 | mm.EXPECT().Size().Return(0) 133 | rm := NewRefCountedMessage(mm, nil) 134 | 135 | mm.EXPECT().Shard().Return(uint32(0)) 136 | require.True(t, rm.Accept(filter)) 137 | 138 | mm.EXPECT().Shard().Return(uint32(1)) 139 | require.False(t, rm.Accept(filter)) 140 | } 141 | 142 | func TestRefCountedMessageOnDropFn(t *testing.T) { 143 | ctrl := gomock.NewController(t) 144 | defer ctrl.Finish() 145 | 146 | mm := NewMockMessage(ctrl) 147 | mm.EXPECT().Size().Return(0) 148 | mm.EXPECT().Finalize(Dropped) 149 | 150 | var called int 151 | fn := func(rm *RefCountedMessage) { 152 | called++ 153 | } 154 | 155 | rm := NewRefCountedMessage(mm, fn) 156 | require.True(t, rm.Drop()) 157 | require.Equal(t, 1, called) 158 | 159 | require.False(t, rm.Drop()) 160 | } 161 | 162 | func TestRefCountedMessageNoBlocking(t *testing.T) { 163 | ctrl := gomock.NewController(t) 164 | defer ctrl.Finish() 165 | 166 | mm := NewMockMessage(ctrl) 167 | mm.EXPECT().Size().Return(0).AnyTimes() 168 | for i := 0; i < 10000; i++ { 169 | rm := NewRefCountedMessage(mm, nil) 170 | var wg sync.WaitGroup 171 | wg.Add(2) 172 | go func() { 173 | rm.IncReads() 174 | rm.IsDroppedOrConsumed() 175 | rm.DecReads() 176 | wg.Done() 177 | }() 178 | go func() { 179 | mm.EXPECT().Finalize(Dropped) 180 | rm.Drop() 181 | wg.Done() 182 | }() 183 | wg.Wait() 184 | } 185 | } 186 | -------------------------------------------------------------------------------- /consumer/options.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2018 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 consumer 22 | 23 | import ( 24 | "time" 25 | 26 | "github.com/m3db/m3msg/protocol/proto" 27 | "github.com/m3db/m3x/instrument" 28 | "github.com/m3db/m3x/pool" 29 | "github.com/m3db/m3x/server" 30 | ) 31 | 32 | var ( 33 | defaultAckBufferSize = 100 34 | defaultAckFlushInterval = time.Second 35 | defaultConnectionBufferSize = 16384 36 | ) 37 | 38 | type options struct { 39 | encOptions proto.Options 40 | decOptions proto.Options 41 | messagePoolOpts pool.ObjectPoolOptions 42 | ackFlushInterval time.Duration 43 | ackBufferSize int 44 | writeBufferSize int 45 | readBufferSize int 46 | iOpts instrument.Options 47 | } 48 | 49 | // NewOptions creates a new options. 50 | func NewOptions() Options { 51 | return &options{ 52 | encOptions: proto.NewOptions(), 53 | decOptions: proto.NewOptions(), 54 | messagePoolOpts: pool.NewObjectPoolOptions(), 55 | ackFlushInterval: defaultAckFlushInterval, 56 | ackBufferSize: defaultAckBufferSize, 57 | writeBufferSize: defaultConnectionBufferSize, 58 | readBufferSize: defaultConnectionBufferSize, 59 | iOpts: instrument.NewOptions(), 60 | } 61 | } 62 | 63 | func (opts *options) EncoderOptions() proto.Options { 64 | return opts.encOptions 65 | } 66 | 67 | func (opts *options) SetEncoderOptions(value proto.Options) Options { 68 | o := *opts 69 | o.encOptions = value 70 | return &o 71 | } 72 | 73 | func (opts *options) DecoderOptions() proto.Options { 74 | return opts.decOptions 75 | } 76 | 77 | func (opts *options) SetDecoderOptions(value proto.Options) Options { 78 | o := *opts 79 | o.decOptions = value 80 | return &o 81 | } 82 | 83 | func (opts *options) MessagePoolOptions() pool.ObjectPoolOptions { 84 | return opts.messagePoolOpts 85 | } 86 | 87 | func (opts *options) SetMessagePoolOptions(value pool.ObjectPoolOptions) Options { 88 | o := *opts 89 | o.messagePoolOpts = value 90 | return &o 91 | } 92 | 93 | func (opts *options) AckFlushInterval() time.Duration { 94 | return opts.ackFlushInterval 95 | } 96 | 97 | func (opts *options) SetAckFlushInterval(value time.Duration) Options { 98 | o := *opts 99 | o.ackFlushInterval = value 100 | return &o 101 | } 102 | 103 | func (opts *options) AckBufferSize() int { 104 | return opts.ackBufferSize 105 | } 106 | 107 | func (opts *options) SetAckBufferSize(value int) Options { 108 | o := *opts 109 | o.ackBufferSize = value 110 | return &o 111 | } 112 | 113 | func (opts *options) ConnectionWriteBufferSize() int { 114 | return opts.writeBufferSize 115 | } 116 | 117 | func (opts *options) SetConnectionWriteBufferSize(value int) Options { 118 | o := *opts 119 | o.writeBufferSize = value 120 | return &o 121 | } 122 | 123 | func (opts *options) ConnectionReadBufferSize() int { 124 | return opts.readBufferSize 125 | } 126 | 127 | func (opts *options) SetConnectionReadBufferSize(value int) Options { 128 | o := *opts 129 | o.readBufferSize = value 130 | return &o 131 | } 132 | 133 | func (opts *options) InstrumentOptions() instrument.Options { 134 | return opts.iOpts 135 | } 136 | 137 | func (opts *options) SetInstrumentOptions(value instrument.Options) Options { 138 | o := *opts 139 | o.iOpts = value 140 | return &o 141 | } 142 | 143 | type serverOptions struct { 144 | consumeFn ConsumeFn 145 | sOpts server.Options 146 | cOpts Options 147 | } 148 | 149 | // NewServerOptions creates ServerOptions. 150 | func NewServerOptions() ServerOptions { 151 | return &serverOptions{ 152 | sOpts: server.NewOptions(), 153 | cOpts: NewOptions(), 154 | } 155 | } 156 | 157 | func (opts *serverOptions) ConsumeFn() ConsumeFn { 158 | return opts.consumeFn 159 | } 160 | 161 | func (opts *serverOptions) SetConsumeFn(value ConsumeFn) ServerOptions { 162 | o := *opts 163 | o.consumeFn = value 164 | return &o 165 | } 166 | 167 | func (opts *serverOptions) ServerOptions() server.Options { 168 | return opts.sOpts 169 | } 170 | 171 | func (opts *serverOptions) SetServerOptions(value server.Options) ServerOptions { 172 | o := *opts 173 | o.sOpts = value 174 | return &o 175 | } 176 | 177 | func (opts *serverOptions) ConsumerOptions() Options { 178 | return opts.cOpts 179 | } 180 | 181 | func (opts *serverOptions) SetConsumerOptions(value Options) ServerOptions { 182 | o := *opts 183 | o.cOpts = value 184 | return &o 185 | } 186 | -------------------------------------------------------------------------------- /topic/types.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2018 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 topic 22 | 23 | import ( 24 | "github.com/m3db/m3cluster/client" 25 | "github.com/m3db/m3cluster/kv" 26 | "github.com/m3db/m3cluster/services" 27 | ) 28 | 29 | // Topic defines the topic of messages. 30 | type Topic interface { 31 | // Name returns the name of the topic. 32 | Name() string 33 | 34 | // SetName sets the name of the topic. 35 | SetName(value string) Topic 36 | 37 | // NumberOfShards returns the total number of shards of the topic. 38 | NumberOfShards() uint32 39 | 40 | // SetNumberOfShards sets the total number of shards of the topic. 41 | SetNumberOfShards(value uint32) Topic 42 | 43 | // ConsumerServices returns the consumers of the topic. 44 | ConsumerServices() []ConsumerService 45 | 46 | // SetConsumerServices sets the consumers of the topic. 47 | SetConsumerServices(value []ConsumerService) Topic 48 | 49 | // Version returns the version of the topic. 50 | Version() int 51 | 52 | // SetVersion sets the version of the topic. 53 | SetVersion(value int) Topic 54 | 55 | // AddConsumerService adds a consumer to the topic. 56 | AddConsumerService(value ConsumerService) (Topic, error) 57 | 58 | // RemoveConsumerService removes a consumer from the topic. 59 | RemoveConsumerService(value services.ServiceID) (Topic, error) 60 | 61 | // UpdateConsumerService updates a consumer in the topic. 62 | UpdateConsumerService(value ConsumerService) (Topic, error) 63 | 64 | // String returns the string representation of the topic. 65 | String() string 66 | 67 | // Validate() validates the topic. 68 | Validate() error 69 | } 70 | 71 | // ConsumerService is a service that consumes the messages in a topic. 72 | type ConsumerService interface { 73 | // ServiceID returns the service id of the consumer service. 74 | ServiceID() services.ServiceID 75 | 76 | // SetServiceID sets the service id of the consumer service. 77 | SetServiceID(value services.ServiceID) ConsumerService 78 | 79 | // ConsumptionType returns the consumption type of the consumer service. 80 | ConsumptionType() ConsumptionType 81 | 82 | // SetConsumptionType sets the consumption type of the consumer service. 83 | SetConsumptionType(value ConsumptionType) ConsumerService 84 | 85 | // MessageTTLNanos returns ttl for each message in nanoseconds. 86 | MessageTTLNanos() int64 87 | 88 | // SetMessageTTLNanos sets ttl for each message in nanoseconds. 89 | SetMessageTTLNanos(value int64) ConsumerService 90 | 91 | // String returns the string representation of the consumer service. 92 | String() string 93 | } 94 | 95 | // Watch watches the updates of a topic. 96 | type Watch interface { 97 | // C returns the notification channel. 98 | C() <-chan struct{} 99 | 100 | // Get returns the latest version of the topic. 101 | Get() (Topic, error) 102 | 103 | // Close stops watching for topic updates. 104 | Close() 105 | } 106 | 107 | // Service provides accessibility to topics. 108 | type Service interface { 109 | // Get returns the topic and version for the given name. 110 | Get(name string) (Topic, error) 111 | 112 | // CheckAndSet sets the topic for the name if the version matches. 113 | CheckAndSet(t Topic, version int) (Topic, error) 114 | 115 | // Delete deletes the topic with the name. 116 | Delete(name string) error 117 | 118 | // Watch returns a topic watch. 119 | Watch(name string) (Watch, error) 120 | } 121 | 122 | // ServiceOptions configures the topic service. 123 | type ServiceOptions interface { 124 | // ConfigService returns the client of config service. 125 | ConfigService() client.Client 126 | 127 | // SetConfigService sets the client of config service. 128 | SetConfigService(c client.Client) ServiceOptions 129 | 130 | // KVOverrideOptions returns the override options for KV store. 131 | KVOverrideOptions() kv.OverrideOptions 132 | 133 | // SetKVOverrideOptions sets the override options for KV store. 134 | SetKVOverrideOptions(value kv.OverrideOptions) ServiceOptions 135 | } 136 | 137 | // ConsumptionType defines how the consumer consumes messages. 138 | type ConsumptionType string 139 | 140 | const ( 141 | // Unknown is the unknown consumption type. 142 | Unknown ConsumptionType = "unknown" 143 | 144 | // Shared means the messages for each shard will be 145 | // shared by all the responsible instances. 146 | Shared ConsumptionType = "shared" 147 | 148 | // Replicated means the messages for each shard will be 149 | // replicated to all the responsible instances. 150 | Replicated ConsumptionType = "replicated" 151 | ) 152 | -------------------------------------------------------------------------------- /producer/buffer/options.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2018 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 buffer 22 | 23 | import ( 24 | "errors" 25 | "time" 26 | 27 | "github.com/m3db/m3x/instrument" 28 | "github.com/m3db/m3x/retry" 29 | ) 30 | 31 | const ( 32 | defaultMaxBufferSize = 100 * 1024 * 1024 // 100MB. 33 | defaultMaxMessageSize = 1 * 1024 * 1024 // 1MB. 34 | defaultCloseCheckInterval = time.Second 35 | defaultDropOldestInterval = time.Second 36 | defaultScanBatchSize = 16 37 | defaultCleanupInitialBackoff = 10 * time.Second 38 | defaultAllowedSpilloverRatio = 0.2 39 | defaultCleanupMaxBackoff = time.Minute 40 | ) 41 | 42 | var ( 43 | errInvalidScanBatchSize = errors.New("invalid scan batch size") 44 | errInvalidMaxMessageSize = errors.New("invalid max message size") 45 | errNegativeMaxBufferSize = errors.New("negative max buffer size") 46 | errNegativeMaxMessageSize = errors.New("negative max message size") 47 | ) 48 | 49 | type bufferOptions struct { 50 | strategy OnFullStrategy 51 | maxBufferSize int 52 | maxMessageSize int 53 | closeCheckInterval time.Duration 54 | dropOldestInterval time.Duration 55 | scanBatchSize int 56 | allowedSpilloverRatio float64 57 | rOpts retry.Options 58 | iOpts instrument.Options 59 | } 60 | 61 | // NewOptions creates Options. 62 | func NewOptions() Options { 63 | return &bufferOptions{ 64 | strategy: DropOldest, 65 | maxBufferSize: defaultMaxBufferSize, 66 | maxMessageSize: defaultMaxMessageSize, 67 | closeCheckInterval: defaultCloseCheckInterval, 68 | dropOldestInterval: defaultDropOldestInterval, 69 | scanBatchSize: defaultScanBatchSize, 70 | allowedSpilloverRatio: defaultAllowedSpilloverRatio, 71 | rOpts: retry.NewOptions(). 72 | SetInitialBackoff(defaultCleanupInitialBackoff). 73 | SetMaxBackoff(defaultCleanupMaxBackoff). 74 | SetForever(true), 75 | iOpts: instrument.NewOptions(), 76 | } 77 | } 78 | 79 | func (opts *bufferOptions) OnFullStrategy() OnFullStrategy { 80 | return opts.strategy 81 | } 82 | 83 | func (opts *bufferOptions) SetOnFullStrategy(value OnFullStrategy) Options { 84 | o := *opts 85 | o.strategy = value 86 | return &o 87 | } 88 | 89 | func (opts *bufferOptions) MaxMessageSize() int { 90 | return opts.maxMessageSize 91 | } 92 | 93 | func (opts *bufferOptions) SetMaxMessageSize(value int) Options { 94 | o := *opts 95 | o.maxMessageSize = value 96 | return &o 97 | } 98 | 99 | func (opts *bufferOptions) MaxBufferSize() int { 100 | return opts.maxBufferSize 101 | } 102 | 103 | func (opts *bufferOptions) SetMaxBufferSize(value int) Options { 104 | o := *opts 105 | o.maxBufferSize = value 106 | return &o 107 | } 108 | 109 | func (opts *bufferOptions) CloseCheckInterval() time.Duration { 110 | return opts.closeCheckInterval 111 | } 112 | 113 | func (opts *bufferOptions) SetCloseCheckInterval(value time.Duration) Options { 114 | o := *opts 115 | o.closeCheckInterval = value 116 | return &o 117 | } 118 | 119 | func (opts *bufferOptions) DropOldestInterval() time.Duration { 120 | return opts.dropOldestInterval 121 | } 122 | 123 | func (opts *bufferOptions) SetDropOldestInterval(value time.Duration) Options { 124 | o := *opts 125 | o.dropOldestInterval = value 126 | return &o 127 | } 128 | 129 | func (opts *bufferOptions) ScanBatchSize() int { 130 | return opts.scanBatchSize 131 | } 132 | 133 | func (opts *bufferOptions) SetScanBatchSize(value int) Options { 134 | o := *opts 135 | o.scanBatchSize = value 136 | return &o 137 | } 138 | 139 | func (opts *bufferOptions) AllowedSpilloverRatio() float64 { 140 | return opts.allowedSpilloverRatio 141 | } 142 | 143 | func (opts *bufferOptions) SetAllowedSpilloverRatio(value float64) Options { 144 | o := *opts 145 | o.allowedSpilloverRatio = value 146 | return &o 147 | } 148 | 149 | func (opts *bufferOptions) CleanupRetryOptions() retry.Options { 150 | return opts.rOpts 151 | } 152 | 153 | func (opts *bufferOptions) SetCleanupRetryOptions(value retry.Options) Options { 154 | o := *opts 155 | o.rOpts = value 156 | return &o 157 | } 158 | 159 | func (opts *bufferOptions) InstrumentOptions() instrument.Options { 160 | return opts.iOpts 161 | } 162 | 163 | func (opts *bufferOptions) SetInstrumentOptions(value instrument.Options) Options { 164 | o := *opts 165 | o.iOpts = value 166 | return &o 167 | } 168 | 169 | func (opts *bufferOptions) Validate() error { 170 | if opts.ScanBatchSize() <= 0 { 171 | return errInvalidScanBatchSize 172 | } 173 | if opts.MaxBufferSize() <= 0 { 174 | return errNegativeMaxBufferSize 175 | } 176 | if opts.MaxMessageSize() <= 0 { 177 | return errNegativeMaxMessageSize 178 | } 179 | if opts.MaxMessageSize() > opts.MaxBufferSize() { 180 | // Max message size can only be as large as max buffer size. 181 | return errInvalidMaxMessageSize 182 | } 183 | return nil 184 | } 185 | -------------------------------------------------------------------------------- /protocol/proto/roundtrip_test.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2018 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 proto 22 | 23 | import ( 24 | "bytes" 25 | "net" 26 | "testing" 27 | 28 | "github.com/m3db/m3msg/generated/proto/msgpb" 29 | "github.com/m3db/m3x/pool" 30 | 31 | "github.com/stretchr/testify/require" 32 | ) 33 | 34 | func TestBaseEncodeDecodeRoundTripWithoutPool(t *testing.T) { 35 | enc := NewEncoder(NewOptions()).(*encoder) 36 | require.Equal(t, 4, len(enc.buffer)) 37 | require.Equal(t, 4, cap(enc.buffer)) 38 | require.Empty(t, enc.Bytes()) 39 | r := bytes.NewReader(nil) 40 | dec := NewDecoder(r, NewOptions()).(*decoder) 41 | require.Equal(t, 4, len(dec.buffer)) 42 | require.Equal(t, 4, cap(dec.buffer)) 43 | encodeMsg := msgpb.Message{ 44 | Metadata: msgpb.Metadata{ 45 | Shard: 1, 46 | Id: 2, 47 | }, 48 | Value: make([]byte, 80), 49 | } 50 | decodeMsg := msgpb.Message{} 51 | 52 | err := enc.Encode(&encodeMsg) 53 | require.NoError(t, err) 54 | require.Equal(t, sizeEncodingLength+encodeMsg.Size(), len(enc.buffer)) 55 | require.Equal(t, sizeEncodingLength+encodeMsg.Size(), cap(enc.buffer)) 56 | 57 | r.Reset(enc.Bytes()) 58 | require.NoError(t, dec.Decode(&decodeMsg)) 59 | require.Equal(t, sizeEncodingLength+decodeMsg.Size(), len(dec.buffer)) 60 | require.Equal(t, sizeEncodingLength+encodeMsg.Size(), cap(dec.buffer)) 61 | } 62 | 63 | func TestBaseEncodeDecodeRoundTripWithPool(t *testing.T) { 64 | p := getBytesPool(2, []int{2, 8, 100}) 65 | p.Init() 66 | 67 | enc := NewEncoder(NewOptions().SetBytesPool(p)).(*encoder) 68 | require.Equal(t, 8, len(enc.buffer)) 69 | require.Equal(t, 8, cap(enc.buffer)) 70 | 71 | r := bytes.NewReader(nil) 72 | dec := NewDecoder(r, NewOptions().SetBytesPool(p)).(*decoder) 73 | require.Equal(t, 8, len(dec.buffer)) 74 | require.Equal(t, 8, cap(dec.buffer)) 75 | encodeMsg := msgpb.Message{ 76 | Metadata: msgpb.Metadata{ 77 | Shard: 1, 78 | Id: 2, 79 | }, 80 | Value: make([]byte, 80), 81 | } 82 | decodeMsg := msgpb.Message{} 83 | 84 | err := enc.Encode(&encodeMsg) 85 | require.NoError(t, err) 86 | require.Equal(t, 100, len(enc.buffer)) 87 | require.Equal(t, 100, cap(enc.buffer)) 88 | 89 | r.Reset(enc.Bytes()) 90 | require.NoError(t, dec.Decode(&decodeMsg)) 91 | require.Equal(t, 100, len(dec.buffer)) 92 | require.Equal(t, 100, cap(dec.buffer)) 93 | } 94 | 95 | func TestResetReader(t *testing.T) { 96 | enc := NewEncoder(nil) 97 | dec := NewDecoder(bytes.NewReader(nil), nil) 98 | encodeMsg := msgpb.Message{ 99 | Metadata: msgpb.Metadata{ 100 | Shard: 1, 101 | Id: 2, 102 | }, 103 | Value: make([]byte, 200), 104 | } 105 | decodeMsg := msgpb.Message{} 106 | 107 | err := enc.Encode(&encodeMsg) 108 | require.NoError(t, err) 109 | require.Error(t, dec.Decode(&decodeMsg)) 110 | 111 | r2 := bytes.NewReader(enc.Bytes()) 112 | dec.(*decoder).ResetReader(r2) 113 | require.NoError(t, dec.Decode(&decodeMsg)) 114 | } 115 | 116 | func TestEncodeMessageLargerThanMaxSize(t *testing.T) { 117 | opts := NewOptions().SetMaxMessageSize(4) 118 | enc := NewEncoder(opts) 119 | encodeMsg := msgpb.Message{ 120 | Metadata: msgpb.Metadata{ 121 | Shard: 1, 122 | Id: 2, 123 | }, 124 | Value: make([]byte, 10), 125 | } 126 | 127 | err := enc.Encode(&encodeMsg) 128 | require.Error(t, err) 129 | require.Contains(t, err.Error(), "larger than maximum supported size") 130 | } 131 | 132 | func TestDecodeMessageLargerThanMaxSize(t *testing.T) { 133 | enc := NewEncoder(nil) 134 | encodeMsg := msgpb.Message{ 135 | Metadata: msgpb.Metadata{ 136 | Shard: 1, 137 | Id: 2, 138 | }, 139 | Value: make([]byte, 10), 140 | } 141 | 142 | err := enc.Encode(&encodeMsg) 143 | require.NoError(t, err) 144 | 145 | decodeMsg := msgpb.Message{} 146 | opts := NewOptions().SetMaxMessageSize(4) 147 | dec := NewDecoder(bytes.NewReader(enc.Bytes()), opts) 148 | err = dec.Decode(&decodeMsg) 149 | require.Error(t, err) 150 | require.Contains(t, err.Error(), "larger than maximum supported size") 151 | } 152 | 153 | func TestEncodeDecodeRoundTrip(t *testing.T) { 154 | enc := NewEncoder(nil) 155 | dec := NewDecoder(nil, nil) 156 | 157 | clientConn, serverConn := net.Pipe() 158 | dec.ResetReader(serverConn) 159 | 160 | testMsg := msgpb.Message{ 161 | Metadata: msgpb.Metadata{ 162 | Shard: 1, 163 | Id: 2, 164 | }, 165 | Value: make([]byte, 10), 166 | } 167 | go func() { 168 | require.NoError(t, enc.Encode(&testMsg)) 169 | _, err := clientConn.Write(enc.Bytes()) 170 | require.NoError(t, err) 171 | }() 172 | var msg msgpb.Message 173 | require.NoError(t, dec.Decode(&msg)) 174 | require.Equal(t, testMsg, msg) 175 | } 176 | 177 | // nolint: unparam 178 | func getBytesPool(bucketSizes int, bucketCaps []int) pool.BytesPool { 179 | buckets := make([]pool.Bucket, len(bucketCaps)) 180 | for i, cap := range bucketCaps { 181 | buckets[i] = pool.Bucket{ 182 | Count: bucketSizes, 183 | Capacity: cap, 184 | } 185 | } 186 | 187 | return pool.NewBytesPool(buckets, nil) 188 | } 189 | -------------------------------------------------------------------------------- /consumer/consumer.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2018 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 consumer 22 | 23 | import ( 24 | "bufio" 25 | "net" 26 | "sync" 27 | "time" 28 | 29 | "github.com/m3db/m3msg/generated/proto/msgpb" 30 | "github.com/m3db/m3msg/protocol/proto" 31 | 32 | "github.com/uber-go/tally" 33 | ) 34 | 35 | type listener struct { 36 | net.Listener 37 | 38 | opts Options 39 | msgPool *messagePool 40 | m metrics 41 | } 42 | 43 | // NewListener creates a consumer listener. 44 | func NewListener(addr string, opts Options) (Listener, error) { 45 | if opts == nil { 46 | opts = NewOptions() 47 | } 48 | lis, err := net.Listen("tcp", addr) 49 | if err != nil { 50 | return nil, err 51 | } 52 | mPool := newMessagePool(opts.MessagePoolOptions()) 53 | mPool.Init() 54 | return &listener{ 55 | Listener: lis, 56 | opts: opts, 57 | msgPool: mPool, 58 | m: newConsumerMetrics(opts.InstrumentOptions().MetricsScope()), 59 | }, nil 60 | } 61 | 62 | func (l *listener) Accept() (Consumer, error) { 63 | conn, err := l.Listener.Accept() 64 | if err != nil { 65 | return nil, err 66 | } 67 | return newConsumer(conn, l.msgPool, l.opts, l.m), nil 68 | } 69 | 70 | type metrics struct { 71 | messageReceived tally.Counter 72 | messageDecodeError tally.Counter 73 | ackSent tally.Counter 74 | ackEncodeError tally.Counter 75 | ackWriteError tally.Counter 76 | } 77 | 78 | func newConsumerMetrics(scope tally.Scope) metrics { 79 | return metrics{ 80 | messageReceived: scope.Counter("message-received"), 81 | messageDecodeError: scope.Counter("message-decode-error"), 82 | ackSent: scope.Counter("ack-sent"), 83 | ackEncodeError: scope.Counter("ack-encode-error"), 84 | ackWriteError: scope.Counter("ack-write-error"), 85 | } 86 | } 87 | 88 | type consumer struct { 89 | sync.Mutex 90 | 91 | opts Options 92 | mPool *messagePool 93 | encoder proto.Encoder 94 | decoder proto.Decoder 95 | w *bufio.Writer 96 | conn net.Conn 97 | 98 | ackPb msgpb.Ack 99 | closed bool 100 | doneCh chan struct{} 101 | wg sync.WaitGroup 102 | m metrics 103 | } 104 | 105 | func newConsumer( 106 | conn net.Conn, 107 | mPool *messagePool, 108 | opts Options, 109 | m metrics, 110 | ) *consumer { 111 | return &consumer{ 112 | opts: opts, 113 | mPool: mPool, 114 | encoder: proto.NewEncoder(opts.EncoderOptions()), 115 | decoder: proto.NewDecoder( 116 | bufio.NewReaderSize(conn, opts.ConnectionReadBufferSize()), 117 | opts.DecoderOptions(), 118 | ), 119 | w: bufio.NewWriterSize(conn, opts.ConnectionWriteBufferSize()), 120 | conn: conn, 121 | closed: false, 122 | doneCh: make(chan struct{}), 123 | m: m, 124 | } 125 | } 126 | 127 | func (c *consumer) Init() { 128 | c.wg.Add(1) 129 | go func() { 130 | c.ackUntilClose() 131 | c.wg.Done() 132 | }() 133 | } 134 | 135 | func (c *consumer) Message() (Message, error) { 136 | m := c.mPool.Get() 137 | m.reset(c) 138 | if err := c.decoder.Decode(m); err != nil { 139 | c.mPool.Put(m) 140 | c.m.messageDecodeError.Inc(1) 141 | return nil, err 142 | } 143 | c.m.messageReceived.Inc(1) 144 | return m, nil 145 | } 146 | 147 | // This function could be called concurrently if messages are being 148 | // processed concurrently. 149 | func (c *consumer) tryAck(m msgpb.Metadata) { 150 | c.Lock() 151 | if c.closed { 152 | c.Unlock() 153 | return 154 | } 155 | c.ackPb.Metadata = append(c.ackPb.Metadata, m) 156 | ackLen := len(c.ackPb.Metadata) 157 | if ackLen < c.opts.AckBufferSize() { 158 | c.Unlock() 159 | return 160 | } 161 | if err := c.encodeAckWithLock(ackLen); err != nil { 162 | c.conn.Close() 163 | } 164 | c.Unlock() 165 | } 166 | 167 | func (c *consumer) ackUntilClose() { 168 | flushTicker := time.NewTicker(c.opts.AckFlushInterval()) 169 | defer flushTicker.Stop() 170 | 171 | for { 172 | select { 173 | case <-flushTicker.C: 174 | c.tryAckAndFlush() 175 | case <-c.doneCh: 176 | c.tryAckAndFlush() 177 | return 178 | } 179 | } 180 | } 181 | 182 | func (c *consumer) tryAckAndFlush() { 183 | c.Lock() 184 | if ackLen := len(c.ackPb.Metadata); ackLen > 0 { 185 | c.encodeAckWithLock(ackLen) 186 | } 187 | c.w.Flush() 188 | c.Unlock() 189 | } 190 | 191 | func (c *consumer) encodeAckWithLock(ackLen int) error { 192 | err := c.encoder.Encode(&c.ackPb) 193 | c.ackPb.Metadata = c.ackPb.Metadata[:0] 194 | if err != nil { 195 | c.m.ackEncodeError.Inc(1) 196 | return err 197 | } 198 | _, err = c.w.Write(c.encoder.Bytes()) 199 | if err != nil { 200 | c.m.ackWriteError.Inc(1) 201 | return err 202 | } 203 | c.m.ackSent.Inc(int64(ackLen)) 204 | return nil 205 | } 206 | 207 | func (c *consumer) Close() { 208 | c.Lock() 209 | if c.closed { 210 | c.Unlock() 211 | return 212 | } 213 | c.closed = true 214 | c.Unlock() 215 | 216 | close(c.doneCh) 217 | c.wg.Wait() 218 | c.conn.Close() 219 | } 220 | 221 | type message struct { 222 | msgpb.Message 223 | 224 | mPool *messagePool 225 | c *consumer 226 | } 227 | 228 | func newMessage(p *messagePool) *message { 229 | return &message{mPool: p} 230 | } 231 | 232 | func (m *message) Bytes() []byte { 233 | return m.Value 234 | } 235 | 236 | func (m *message) Ack() { 237 | m.c.tryAck(m.Metadata) 238 | if m.mPool != nil { 239 | m.mPool.Put(m) 240 | } 241 | } 242 | 243 | func (m *message) reset(c *consumer) { 244 | m.c = c 245 | resetProto(&m.Message) 246 | } 247 | 248 | func resetProto(m *msgpb.Message) { 249 | m.Metadata.Id = 0 250 | m.Metadata.Shard = 0 251 | m.Value = m.Value[:0] 252 | } 253 | -------------------------------------------------------------------------------- /producer/producer_mock.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2018 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 | // Code generated by MockGen. DO NOT EDIT. 22 | // Source: github.com/m3db/m3msg/producer (interfaces: Message,Producer) 23 | 24 | // Package producer is a generated GoMock package. 25 | package producer 26 | 27 | import ( 28 | "reflect" 29 | 30 | "github.com/m3db/m3cluster/services" 31 | 32 | "github.com/golang/mock/gomock" 33 | ) 34 | 35 | // MockMessage is a mock of Message interface 36 | type MockMessage struct { 37 | ctrl *gomock.Controller 38 | recorder *MockMessageMockRecorder 39 | } 40 | 41 | // MockMessageMockRecorder is the mock recorder for MockMessage 42 | type MockMessageMockRecorder struct { 43 | mock *MockMessage 44 | } 45 | 46 | // NewMockMessage creates a new mock instance 47 | func NewMockMessage(ctrl *gomock.Controller) *MockMessage { 48 | mock := &MockMessage{ctrl: ctrl} 49 | mock.recorder = &MockMessageMockRecorder{mock} 50 | return mock 51 | } 52 | 53 | // EXPECT returns an object that allows the caller to indicate expected use 54 | func (m *MockMessage) EXPECT() *MockMessageMockRecorder { 55 | return m.recorder 56 | } 57 | 58 | // Bytes mocks base method 59 | func (m *MockMessage) Bytes() []byte { 60 | ret := m.ctrl.Call(m, "Bytes") 61 | ret0, _ := ret[0].([]byte) 62 | return ret0 63 | } 64 | 65 | // Bytes indicates an expected call of Bytes 66 | func (mr *MockMessageMockRecorder) Bytes() *gomock.Call { 67 | return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Bytes", reflect.TypeOf((*MockMessage)(nil).Bytes)) 68 | } 69 | 70 | // Finalize mocks base method 71 | func (m *MockMessage) Finalize(arg0 FinalizeReason) { 72 | m.ctrl.Call(m, "Finalize", arg0) 73 | } 74 | 75 | // Finalize indicates an expected call of Finalize 76 | func (mr *MockMessageMockRecorder) Finalize(arg0 interface{}) *gomock.Call { 77 | return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Finalize", reflect.TypeOf((*MockMessage)(nil).Finalize), arg0) 78 | } 79 | 80 | // Shard mocks base method 81 | func (m *MockMessage) Shard() uint32 { 82 | ret := m.ctrl.Call(m, "Shard") 83 | ret0, _ := ret[0].(uint32) 84 | return ret0 85 | } 86 | 87 | // Shard indicates an expected call of Shard 88 | func (mr *MockMessageMockRecorder) Shard() *gomock.Call { 89 | return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Shard", reflect.TypeOf((*MockMessage)(nil).Shard)) 90 | } 91 | 92 | // Size mocks base method 93 | func (m *MockMessage) Size() int { 94 | ret := m.ctrl.Call(m, "Size") 95 | ret0, _ := ret[0].(int) 96 | return ret0 97 | } 98 | 99 | // Size indicates an expected call of Size 100 | func (mr *MockMessageMockRecorder) Size() *gomock.Call { 101 | return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Size", reflect.TypeOf((*MockMessage)(nil).Size)) 102 | } 103 | 104 | // MockProducer is a mock of Producer interface 105 | type MockProducer struct { 106 | ctrl *gomock.Controller 107 | recorder *MockProducerMockRecorder 108 | } 109 | 110 | // MockProducerMockRecorder is the mock recorder for MockProducer 111 | type MockProducerMockRecorder struct { 112 | mock *MockProducer 113 | } 114 | 115 | // NewMockProducer creates a new mock instance 116 | func NewMockProducer(ctrl *gomock.Controller) *MockProducer { 117 | mock := &MockProducer{ctrl: ctrl} 118 | mock.recorder = &MockProducerMockRecorder{mock} 119 | return mock 120 | } 121 | 122 | // EXPECT returns an object that allows the caller to indicate expected use 123 | func (m *MockProducer) EXPECT() *MockProducerMockRecorder { 124 | return m.recorder 125 | } 126 | 127 | // Close mocks base method 128 | func (m *MockProducer) Close(arg0 CloseType) { 129 | m.ctrl.Call(m, "Close", arg0) 130 | } 131 | 132 | // Close indicates an expected call of Close 133 | func (mr *MockProducerMockRecorder) Close(arg0 interface{}) *gomock.Call { 134 | return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Close", reflect.TypeOf((*MockProducer)(nil).Close), arg0) 135 | } 136 | 137 | // Init mocks base method 138 | func (m *MockProducer) Init() error { 139 | ret := m.ctrl.Call(m, "Init") 140 | ret0, _ := ret[0].(error) 141 | return ret0 142 | } 143 | 144 | // Init indicates an expected call of Init 145 | func (mr *MockProducerMockRecorder) Init() *gomock.Call { 146 | return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Init", reflect.TypeOf((*MockProducer)(nil).Init)) 147 | } 148 | 149 | // Produce mocks base method 150 | func (m *MockProducer) Produce(arg0 Message) error { 151 | ret := m.ctrl.Call(m, "Produce", arg0) 152 | ret0, _ := ret[0].(error) 153 | return ret0 154 | } 155 | 156 | // Produce indicates an expected call of Produce 157 | func (mr *MockProducerMockRecorder) Produce(arg0 interface{}) *gomock.Call { 158 | return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Produce", reflect.TypeOf((*MockProducer)(nil).Produce), arg0) 159 | } 160 | 161 | // RegisterFilter mocks base method 162 | func (m *MockProducer) RegisterFilter(arg0 services.ServiceID, arg1 FilterFunc) { 163 | m.ctrl.Call(m, "RegisterFilter", arg0, arg1) 164 | } 165 | 166 | // RegisterFilter indicates an expected call of RegisterFilter 167 | func (mr *MockProducerMockRecorder) RegisterFilter(arg0, arg1 interface{}) *gomock.Call { 168 | return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "RegisterFilter", reflect.TypeOf((*MockProducer)(nil).RegisterFilter), arg0, arg1) 169 | } 170 | 171 | // UnregisterFilter mocks base method 172 | func (m *MockProducer) UnregisterFilter(arg0 services.ServiceID) { 173 | m.ctrl.Call(m, "UnregisterFilter", arg0) 174 | } 175 | 176 | // UnregisterFilter indicates an expected call of UnregisterFilter 177 | func (mr *MockProducerMockRecorder) UnregisterFilter(arg0 interface{}) *gomock.Call { 178 | return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UnregisterFilter", reflect.TypeOf((*MockProducer)(nil).UnregisterFilter), arg0) 179 | } 180 | -------------------------------------------------------------------------------- /producer/config/writer.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2018 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 config 22 | 23 | import ( 24 | "time" 25 | 26 | "github.com/m3db/m3cluster/client" 27 | "github.com/m3db/m3cluster/kv" 28 | "github.com/m3db/m3cluster/services" 29 | "github.com/m3db/m3msg/producer/writer" 30 | "github.com/m3db/m3msg/protocol/proto" 31 | "github.com/m3db/m3msg/topic" 32 | "github.com/m3db/m3x/instrument" 33 | "github.com/m3db/m3x/pool" 34 | "github.com/m3db/m3x/retry" 35 | ) 36 | 37 | // ConnectionConfiguration configs the connection options. 38 | type ConnectionConfiguration struct { 39 | DialTimeout *time.Duration `yaml:"dialTimeout"` 40 | WriteTimeout *time.Duration `yaml:"writeTimeout"` 41 | KeepAlivePeriod *time.Duration `yaml:"keepAlivePeriod"` 42 | ResetDelay *time.Duration `yaml:"resetDelay"` 43 | Retry *retry.Configuration `yaml:"retry"` 44 | FlushInterval *time.Duration `yaml:"flushInterval"` 45 | WriteBufferSize *int `yaml:"writeBufferSize"` 46 | ReadBufferSize *int `yaml:"readBufferSize"` 47 | } 48 | 49 | // NewOptions creates connection options. 50 | func (c *ConnectionConfiguration) NewOptions(iOpts instrument.Options) writer.ConnectionOptions { 51 | opts := writer.NewConnectionOptions() 52 | if c.DialTimeout != nil { 53 | opts = opts.SetDialTimeout(*c.DialTimeout) 54 | } 55 | if c.WriteTimeout != nil { 56 | opts = opts.SetWriteTimeout(*c.WriteTimeout) 57 | } 58 | if c.KeepAlivePeriod != nil { 59 | opts = opts.SetKeepAlivePeriod(*c.KeepAlivePeriod) 60 | } 61 | if c.ResetDelay != nil { 62 | opts = opts.SetResetDelay(*c.ResetDelay) 63 | } 64 | if c.Retry != nil { 65 | opts = opts.SetRetryOptions(c.Retry.NewOptions(iOpts.MetricsScope())) 66 | } 67 | if c.FlushInterval != nil { 68 | opts = opts.SetFlushInterval(*c.FlushInterval) 69 | } 70 | if c.WriteBufferSize != nil { 71 | opts = opts.SetWriteBufferSize(*c.WriteBufferSize) 72 | } 73 | if c.ReadBufferSize != nil { 74 | opts = opts.SetReadBufferSize(*c.ReadBufferSize) 75 | } 76 | return opts.SetInstrumentOptions(iOpts) 77 | } 78 | 79 | // WriterConfiguration configs the writer options. 80 | type WriterConfiguration struct { 81 | TopicName string `yaml:"topicName" validate:"nonzero"` 82 | TopicServiceOverride kv.OverrideConfiguration `yaml:"topicServiceOverride"` 83 | TopicWatchInitTimeout *time.Duration `yaml:"topicWatchInitTimeout"` 84 | PlacementServiceOverride services.OverrideConfiguration `yaml:"placementServiceOverride"` 85 | PlacementWatchInitTimeout *time.Duration `yaml:"placementWatchInitTimeout"` 86 | MessagePool *pool.ObjectPoolConfiguration `yaml:"messagePool"` 87 | MessageRetry *retry.Configuration `yaml:"messageRetry"` 88 | MessageQueueNewWritesScanInterval *time.Duration `yaml:"messageQueueNewWritesScanInterval"` 89 | MessageQueueFullScanInterval *time.Duration `yaml:"messageQueueFullScanInterval"` 90 | MessageQueueScanBatchSize *int `yaml:"messageQueueScanBatchSize"` 91 | InitialAckMapSize *int `yaml:"initialAckMapSize"` 92 | CloseCheckInterval *time.Duration `yaml:"closeCheckInterval"` 93 | AckErrorRetry *retry.Configuration `yaml:"ackErrorRetry"` 94 | Encoder *proto.Configuration `yaml:"encoder"` 95 | Decoder *proto.Configuration `yaml:"decoder"` 96 | Connection *ConnectionConfiguration `yaml:"connection"` 97 | } 98 | 99 | // NewOptions creates writer options. 100 | func (c *WriterConfiguration) NewOptions( 101 | cs client.Client, 102 | iOpts instrument.Options, 103 | ) (writer.Options, error) { 104 | opts := writer.NewOptions().SetTopicName(c.TopicName) 105 | kvOpts, err := c.TopicServiceOverride.NewOverrideOptions() 106 | if err != nil { 107 | return nil, err 108 | } 109 | ts, err := topic.NewService( 110 | topic.NewServiceOptions(). 111 | SetConfigService(cs). 112 | SetKVOverrideOptions(kvOpts), 113 | ) 114 | if err != nil { 115 | return nil, err 116 | } 117 | opts = opts.SetTopicService(ts) 118 | if c.TopicWatchInitTimeout != nil { 119 | opts = opts.SetTopicWatchInitTimeout(*c.TopicWatchInitTimeout) 120 | } 121 | sd, err := cs.Services(c.PlacementServiceOverride.NewOptions()) 122 | if err != nil { 123 | return nil, err 124 | } 125 | opts = opts.SetServiceDiscovery(sd) 126 | if c.PlacementWatchInitTimeout != nil { 127 | opts = opts.SetPlacementWatchInitTimeout(*c.PlacementWatchInitTimeout) 128 | } 129 | if c.MessagePool != nil { 130 | opts = opts.SetMessagePoolOptions(c.MessagePool.NewObjectPoolOptions(iOpts)) 131 | } 132 | if c.MessageRetry != nil { 133 | opts = opts.SetMessageRetryOptions(c.MessageRetry.NewOptions(iOpts.MetricsScope())) 134 | } 135 | if c.MessageQueueNewWritesScanInterval != nil { 136 | opts = opts.SetMessageQueueNewWritesScanInterval(*c.MessageQueueNewWritesScanInterval) 137 | } 138 | if c.MessageQueueFullScanInterval != nil { 139 | opts = opts.SetMessageQueueFullScanInterval(*c.MessageQueueFullScanInterval) 140 | } 141 | if c.MessageQueueScanBatchSize != nil { 142 | opts = opts.SetMessageQueueScanBatchSize(*c.MessageQueueScanBatchSize) 143 | } 144 | if c.InitialAckMapSize != nil { 145 | opts = opts.SetInitialAckMapSize(*c.InitialAckMapSize) 146 | } 147 | if c.CloseCheckInterval != nil { 148 | opts = opts.SetCloseCheckInterval(*c.CloseCheckInterval) 149 | } 150 | if c.AckErrorRetry != nil { 151 | opts = opts.SetAckErrorRetryOptions(c.AckErrorRetry.NewOptions(iOpts.MetricsScope())) 152 | } 153 | if c.Encoder != nil { 154 | opts = opts.SetEncoderOptions(c.Encoder.NewOptions(iOpts)) 155 | } 156 | if c.Decoder != nil { 157 | opts = opts.SetDecoderOptions(c.Decoder.NewOptions(iOpts)) 158 | } 159 | if c.Connection != nil { 160 | opts = opts.SetConnectionOptions(c.Connection.NewOptions(iOpts)) 161 | } 162 | return opts.SetInstrumentOptions(iOpts), nil 163 | } 164 | -------------------------------------------------------------------------------- /topic/topic_test.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2018 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 topic 22 | 23 | import ( 24 | "testing" 25 | "time" 26 | 27 | "github.com/m3db/m3cluster/services" 28 | 29 | "github.com/stretchr/testify/require" 30 | ) 31 | 32 | func TestTopicAddConsumer(t *testing.T) { 33 | cs1 := NewConsumerService(). 34 | SetConsumptionType(Shared). 35 | SetServiceID(services.NewServiceID(). 36 | SetName("s1"). 37 | SetEnvironment("env1"). 38 | SetZone("zone1"), 39 | ) 40 | cs2 := NewConsumerService(). 41 | SetConsumptionType(Shared). 42 | SetServiceID(services.NewServiceID(). 43 | SetName("s2"). 44 | SetEnvironment("env2"). 45 | SetZone("zone2"), 46 | ) 47 | tpc := NewTopic(). 48 | SetName("testName"). 49 | SetNumberOfShards(1024). 50 | SetVersion(5). 51 | SetConsumerServices( 52 | []ConsumerService{cs1}, 53 | ) 54 | 55 | _, err := tpc.AddConsumerService( 56 | NewConsumerService(). 57 | SetConsumptionType(Shared). 58 | SetServiceID(services.NewServiceID(). 59 | SetName("s1"). 60 | SetEnvironment("env1"). 61 | SetZone("zone1"), 62 | ), 63 | ) 64 | require.Error(t, err) 65 | require.Contains(t, err.Error(), cs1.ServiceID().String()) 66 | 67 | tpc, err = tpc.AddConsumerService(cs2) 68 | require.NoError(t, err) 69 | require.Equal(t, []ConsumerService{cs1, cs2}, tpc.ConsumerServices()) 70 | } 71 | 72 | func TestTopicRemoveConsumer(t *testing.T) { 73 | cs1 := NewConsumerService(). 74 | SetConsumptionType(Shared). 75 | SetServiceID(services.NewServiceID(). 76 | SetName("s1"). 77 | SetEnvironment("env1"). 78 | SetZone("zone1"), 79 | ) 80 | cs2 := NewConsumerService(). 81 | SetConsumptionType(Shared). 82 | SetServiceID(services.NewServiceID(). 83 | SetName("s2"). 84 | SetEnvironment("env2"). 85 | SetZone("zone2"), 86 | ) 87 | tpc := NewTopic(). 88 | SetName("testName"). 89 | SetNumberOfShards(1024). 90 | SetVersion(5). 91 | SetConsumerServices( 92 | []ConsumerService{cs1}, 93 | ) 94 | 95 | _, err := tpc.RemoveConsumerService(cs2.ServiceID()) 96 | require.Error(t, err) 97 | require.Contains(t, err.Error(), cs2.ServiceID().String()) 98 | 99 | tpc, err = tpc.RemoveConsumerService( 100 | services.NewServiceID(). 101 | SetName("s1"). 102 | SetEnvironment("env1"). 103 | SetZone("zone1"), 104 | ) 105 | require.NoError(t, err) 106 | require.Empty(t, tpc.ConsumerServices()) 107 | } 108 | 109 | func TestTopicUpdateConsumer(t *testing.T) { 110 | cs1 := NewConsumerService(). 111 | SetConsumptionType(Shared). 112 | SetServiceID(services.NewServiceID(). 113 | SetName("s1"). 114 | SetEnvironment("env1"). 115 | SetZone("zone1"), 116 | ) 117 | tpc := NewTopic(). 118 | SetName("testName"). 119 | SetNumberOfShards(1024). 120 | SetConsumerServices( 121 | []ConsumerService{cs1}, 122 | ) 123 | 124 | _, err := tpc.UpdateConsumerService(cs1.SetConsumptionType(Replicated)) 125 | require.Error(t, err) 126 | require.Contains(t, err.Error(), "could not change consumption type") 127 | 128 | _, err = tpc.UpdateConsumerService(cs1.SetServiceID(services.NewServiceID().SetName("foo"))) 129 | require.Error(t, err) 130 | require.Contains(t, err.Error(), "could not find consumer service") 131 | 132 | require.Equal(t, []ConsumerService{cs1}, tpc.ConsumerServices()) 133 | cs2 := NewConsumerService(). 134 | SetConsumptionType(Shared). 135 | SetServiceID(services.NewServiceID(). 136 | SetName("s1"). 137 | SetEnvironment("env1"). 138 | SetZone("zone1"), 139 | ).SetMessageTTLNanos(500) 140 | tpc, err = tpc.UpdateConsumerService(cs2) 141 | require.NoError(t, err) 142 | require.Equal(t, []ConsumerService{cs2}, tpc.ConsumerServices()) 143 | require.Equal(t, int64(0), cs1.MessageTTLNanos()) 144 | require.Equal(t, int64(500), cs2.MessageTTLNanos()) 145 | } 146 | 147 | func TestTopicString(t *testing.T) { 148 | cs1 := NewConsumerService(). 149 | SetConsumptionType(Shared). 150 | SetServiceID(services.NewServiceID(). 151 | SetName("s1"). 152 | SetEnvironment("env1"). 153 | SetZone("zone1"), 154 | ) 155 | cs2 := NewConsumerService(). 156 | SetConsumptionType(Shared). 157 | SetServiceID(services.NewServiceID(). 158 | SetName("s2"). 159 | SetEnvironment("env2"). 160 | SetZone("zone2"), 161 | ). 162 | SetMessageTTLNanos(int64(time.Minute)) 163 | tpc := NewTopic(). 164 | SetName("testName"). 165 | SetNumberOfShards(1024). 166 | SetVersion(5). 167 | SetConsumerServices( 168 | []ConsumerService{cs1, cs2}, 169 | ) 170 | str := ` 171 | { 172 | version: 5 173 | name: testName 174 | numOfShards: 1024 175 | consumerServices: { 176 | {service: [name: s1, env: env1, zone: zone1], consumption type: shared} 177 | {service: [name: s2, env: env2, zone: zone2], consumption type: shared, ttl: 1m0s} 178 | } 179 | } 180 | ` 181 | require.Equal(t, str, tpc.String()) 182 | } 183 | 184 | func TestTopicValidation(t *testing.T) { 185 | topic := NewTopic() 186 | err := topic.Validate() 187 | require.Error(t, err) 188 | require.Equal(t, errEmptyName, err) 189 | 190 | topic = topic.SetName("name") 191 | err = topic.Validate() 192 | require.Error(t, err) 193 | require.Equal(t, errZeroShards, err) 194 | 195 | topic = topic.SetNumberOfShards(1024) 196 | err = topic.Validate() 197 | require.NoError(t, err) 198 | 199 | cs1 := NewConsumerService(). 200 | SetConsumptionType(Shared). 201 | SetServiceID(services.NewServiceID(). 202 | SetName("s1"). 203 | SetEnvironment("env1"). 204 | SetZone("zone1"), 205 | ) 206 | topic = topic.SetConsumerServices([]ConsumerService{ 207 | cs1, cs1, 208 | }) 209 | err = topic.Validate() 210 | require.Error(t, err) 211 | require.Contains(t, err.Error(), "duplicated consumer") 212 | 213 | topic = topic.SetConsumerServices([]ConsumerService{ 214 | cs1, cs1.SetConsumptionType(Replicated), 215 | }) 216 | err = topic.Validate() 217 | require.Error(t, err) 218 | require.Contains(t, err.Error(), "duplicated consumer") 219 | 220 | topic = topic.SetConsumerServices([]ConsumerService{ 221 | cs1, 222 | }) 223 | err = topic.Validate() 224 | require.NoError(t, err) 225 | } 226 | 227 | func TestConsumerService(t *testing.T) { 228 | sid := services.NewServiceID().SetName("s").SetEnvironment("env").SetZone("zone") 229 | cs := NewConsumerService().SetConsumptionType(Shared).SetServiceID(sid).SetMessageTTLNanos(int64(time.Second)) 230 | require.Equal(t, sid, cs.ServiceID()) 231 | require.Equal(t, Shared, cs.ConsumptionType()) 232 | require.Equal(t, int64(time.Second), cs.MessageTTLNanos()) 233 | require.Equal(t, "{service: [name: s, env: env, zone: zone], consumption type: shared, ttl: 1s}", cs.String()) 234 | } 235 | -------------------------------------------------------------------------------- /producer/writer/writer.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2018 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 writer 22 | 23 | import ( 24 | "errors" 25 | "fmt" 26 | "sync" 27 | 28 | "github.com/m3db/m3cluster/services" 29 | "github.com/m3db/m3msg/producer" 30 | "github.com/m3db/m3msg/topic" 31 | xerrors "github.com/m3db/m3x/errors" 32 | "github.com/m3db/m3x/log" 33 | "github.com/m3db/m3x/watch" 34 | 35 | "github.com/uber-go/tally" 36 | ) 37 | 38 | var ( 39 | errWriterClosed = errors.New("writer is closed") 40 | ) 41 | 42 | type writerMetrics struct { 43 | topicUpdateSuccess tally.Counter 44 | topicUpdateError tally.Counter 45 | invalidTopicUpdate tally.Counter 46 | invalidShard tally.Counter 47 | numConsumerServices tally.Gauge 48 | } 49 | 50 | func newWriterMetrics(scope tally.Scope) writerMetrics { 51 | return writerMetrics{ 52 | topicUpdateSuccess: scope.Counter("topic-update-success"), 53 | topicUpdateError: scope.Counter("topic-update-error"), 54 | invalidTopicUpdate: scope.Counter("invalid-topic"), 55 | invalidShard: scope.Tagged(map[string]string{"reason": "invalid-shard"}). 56 | Counter("invalid-write"), 57 | numConsumerServices: scope.Gauge("num-consumer-services"), 58 | } 59 | } 60 | 61 | type writer struct { 62 | sync.RWMutex 63 | 64 | topic string 65 | ts topic.Service 66 | opts Options 67 | logger log.Logger 68 | 69 | value watch.Value 70 | initType initType 71 | numShards uint32 72 | consumerServiceWriters map[string]consumerServiceWriter 73 | filterRegistry map[string]producer.FilterFunc 74 | isClosed bool 75 | m writerMetrics 76 | 77 | processFn watch.ProcessFn 78 | } 79 | 80 | // NewWriter creates a new writer. 81 | func NewWriter(opts Options) producer.Writer { 82 | w := &writer{ 83 | topic: opts.TopicName(), 84 | ts: opts.TopicService(), 85 | opts: opts, 86 | logger: opts.InstrumentOptions().Logger(), 87 | initType: failOnError, 88 | consumerServiceWriters: make(map[string]consumerServiceWriter), 89 | filterRegistry: make(map[string]producer.FilterFunc), 90 | isClosed: false, 91 | m: newWriterMetrics(opts.InstrumentOptions().MetricsScope()), 92 | } 93 | w.processFn = w.process 94 | return w 95 | } 96 | 97 | func (w *writer) Write(rm *producer.RefCountedMessage) error { 98 | w.RLock() 99 | if w.isClosed { 100 | rm.Drop() 101 | w.RUnlock() 102 | return errWriterClosed 103 | } 104 | shard := rm.Shard() 105 | if shard >= w.numShards { 106 | w.m.invalidShard.Inc(1) 107 | rm.Drop() 108 | w.RUnlock() 109 | return fmt.Errorf("could not write message for shard %d which is larger than max shard id %d", shard, w.numShards-1) 110 | } 111 | // NB(cw): Need to inc ref here in case a consumer service 112 | // writes the message too fast and close the message. 113 | rm.IncRef() 114 | for _, csw := range w.consumerServiceWriters { 115 | csw.Write(rm) 116 | } 117 | rm.DecRef() 118 | w.RUnlock() 119 | return nil 120 | } 121 | 122 | func (w *writer) Init() error { 123 | newUpdatableFn := func() (watch.Updatable, error) { 124 | return w.ts.Watch(w.topic) 125 | } 126 | getUpdateFn := func(value watch.Updatable) (interface{}, error) { 127 | t, err := value.(topic.Watch).Get() 128 | if err != nil { 129 | w.m.invalidTopicUpdate.Inc(1) 130 | return nil, err 131 | } 132 | return t, nil 133 | } 134 | vOptions := watch.NewOptions(). 135 | SetInitWatchTimeout(w.opts.TopicWatchInitTimeout()). 136 | SetInstrumentOptions(w.opts.InstrumentOptions()). 137 | SetNewUpdatableFn(newUpdatableFn). 138 | SetGetUpdateFn(getUpdateFn). 139 | SetProcessFn(w.processFn) 140 | w.value = watch.NewValue(vOptions) 141 | if err := w.value.Watch(); err != nil { 142 | return fmt.Errorf("writer init error: %v", err) 143 | } 144 | return nil 145 | } 146 | 147 | func (w *writer) process(update interface{}) error { 148 | t := update.(topic.Topic) 149 | if err := t.Validate(); err != nil { 150 | return err 151 | } 152 | // We don't allow changing number of shards for topics, it will be 153 | // prevented on topic service side, but also being defensive here as well. 154 | if w.numShards != 0 && w.numShards != t.NumberOfShards() { 155 | w.m.topicUpdateError.Inc(1) 156 | return fmt.Errorf("invalid topic update with %d shards, expecting %d", t.NumberOfShards(), w.numShards) 157 | } 158 | var ( 159 | iOpts = w.opts.InstrumentOptions() 160 | newConsumerServiceWriters = make(map[string]consumerServiceWriter, len(t.ConsumerServices())) 161 | toBeClosed []consumerServiceWriter 162 | multiErr xerrors.MultiError 163 | ) 164 | for _, cs := range t.ConsumerServices() { 165 | key := cs.ServiceID().String() 166 | csw, ok := w.consumerServiceWriters[key] 167 | if ok { 168 | csw.SetMessageTTLNanos(cs.MessageTTLNanos()) 169 | newConsumerServiceWriters[key] = csw 170 | continue 171 | } 172 | scope := iOpts.MetricsScope().Tagged(map[string]string{ 173 | "consumer-service-name": cs.ServiceID().Name(), 174 | "consumer-service-zone": cs.ServiceID().Zone(), 175 | "consumer-service-env": cs.ServiceID().Environment(), 176 | "consumption-type": cs.ConsumptionType().String(), 177 | }) 178 | csw, err := newConsumerServiceWriter(cs, t.NumberOfShards(), w.opts.SetInstrumentOptions(iOpts.SetMetricsScope(scope))) 179 | if err != nil { 180 | w.logger.Errorf("could not create consumer service writer for %s: %v", cs.String(), err) 181 | multiErr = multiErr.Add(err) 182 | continue 183 | } 184 | if err = csw.Init(w.initType); err != nil { 185 | w.logger.Errorf("could not init consumer service writer for %s: %v", cs.String(), err) 186 | multiErr = multiErr.Add(err) 187 | // Could not initialize the consumer service, simply close it. 188 | csw.Close() 189 | continue 190 | } 191 | csw.SetMessageTTLNanos(cs.MessageTTLNanos()) 192 | newConsumerServiceWriters[key] = csw 193 | w.logger.Infof("initialized consumer service writer for %s", cs.String()) 194 | } 195 | for key, csw := range w.consumerServiceWriters { 196 | if _, ok := newConsumerServiceWriters[key]; !ok { 197 | toBeClosed = append(toBeClosed, csw) 198 | } 199 | } 200 | // Allow InitValueError for any future topic updates after starting up. 201 | // This is to handle the case when a new consumer service got added to 202 | // the topic, but the producer could not get initial value for its 203 | // placement. We will continue to watch for placement updates for the new 204 | // consumer service in the background, so the producer can write to it once 205 | // the placement came in. 206 | w.initType = allowInitValueError 207 | w.m.numConsumerServices.Update(float64(len(newConsumerServiceWriters))) 208 | 209 | // Apply the new consumer service writers. 210 | w.Lock() 211 | for key, csw := range newConsumerServiceWriters { 212 | if filter, ok := w.filterRegistry[key]; ok { 213 | csw.RegisterFilter(filter) 214 | } 215 | } 216 | w.consumerServiceWriters = newConsumerServiceWriters 217 | w.numShards = t.NumberOfShards() 218 | w.Unlock() 219 | 220 | // Close removed consumer service. 221 | go func() { 222 | for _, csw := range toBeClosed { 223 | csw.Close() 224 | } 225 | }() 226 | 227 | if err := multiErr.FinalError(); err != nil { 228 | w.m.topicUpdateError.Inc(1) 229 | return err 230 | } 231 | w.m.topicUpdateSuccess.Inc(1) 232 | return nil 233 | } 234 | 235 | func (w *writer) Close() { 236 | w.Lock() 237 | if w.isClosed { 238 | w.Unlock() 239 | return 240 | } 241 | w.isClosed = true 242 | w.Unlock() 243 | 244 | w.value.Unwatch() 245 | for _, csw := range w.consumerServiceWriters { 246 | csw.Close() 247 | } 248 | } 249 | 250 | func (w *writer) RegisterFilter(sid services.ServiceID, filter producer.FilterFunc) { 251 | w.Lock() 252 | defer w.Unlock() 253 | 254 | key := sid.String() 255 | w.filterRegistry[key] = filter 256 | csw, ok := w.consumerServiceWriters[key] 257 | if ok { 258 | csw.RegisterFilter(filter) 259 | } 260 | } 261 | 262 | func (w *writer) UnregisterFilter(sid services.ServiceID) { 263 | w.Lock() 264 | defer w.Unlock() 265 | 266 | key := sid.String() 267 | delete(w.filterRegistry, key) 268 | csw, ok := w.consumerServiceWriters[key] 269 | if ok { 270 | csw.UnregisterFilter() 271 | } 272 | } 273 | -------------------------------------------------------------------------------- /producer/writer/shard_writer.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2018 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 writer 22 | 23 | import ( 24 | "sync" 25 | 26 | "github.com/m3db/m3cluster/placement" 27 | "github.com/m3db/m3msg/producer" 28 | "github.com/m3db/m3x/log" 29 | 30 | "go.uber.org/atomic" 31 | ) 32 | 33 | type shardWriter interface { 34 | // Write writes the reference counted message, this needs to be thread safe. 35 | Write(rm *producer.RefCountedMessage) 36 | 37 | // UpdateInstances updates the instances responsible for this shard. 38 | UpdateInstances( 39 | instances []placement.Instance, 40 | cws map[string]consumerWriter, 41 | ) 42 | 43 | // SetMessageTTLNanos sets the message ttl nanoseconds. 44 | SetMessageTTLNanos(value int64) 45 | 46 | // Close closes the shard writer. 47 | Close() 48 | 49 | // QueueSize returns the number of messages queued for the shard. 50 | QueueSize() int 51 | } 52 | 53 | type sharedShardWriter struct { 54 | instances map[string]struct{} 55 | mw messageWriter 56 | isClosed *atomic.Bool 57 | } 58 | 59 | func newSharedShardWriter( 60 | shard uint32, 61 | router ackRouter, 62 | mPool messagePool, 63 | opts Options, 64 | m messageWriterMetrics, 65 | ) shardWriter { 66 | replicatedShardID := uint64(shard) 67 | mw := newMessageWriter(replicatedShardID, mPool, opts, m) 68 | mw.Init() 69 | router.Register(replicatedShardID, mw) 70 | return &sharedShardWriter{ 71 | instances: make(map[string]struct{}), 72 | mw: mw, 73 | isClosed: atomic.NewBool(false), 74 | } 75 | } 76 | 77 | func (w *sharedShardWriter) Write(rm *producer.RefCountedMessage) { 78 | w.mw.Write(rm) 79 | } 80 | 81 | // This is not thread safe, must be called in one thread. 82 | func (w *sharedShardWriter) UpdateInstances( 83 | instances []placement.Instance, 84 | cws map[string]consumerWriter, 85 | ) { 86 | var ( 87 | newInstancesMap = make(map[string]struct{}, len(instances)) 88 | toBeDeleted = w.instances 89 | ) 90 | for _, instance := range instances { 91 | id := instance.Endpoint() 92 | newInstancesMap[id] = struct{}{} 93 | if _, ok := toBeDeleted[id]; ok { 94 | // Existing instance. 95 | delete(toBeDeleted, id) 96 | continue 97 | } 98 | // Add the consumer writer to the message writer. 99 | w.mw.AddConsumerWriter(cws[id]) 100 | } 101 | for id := range toBeDeleted { 102 | w.mw.RemoveConsumerWriter(id) 103 | } 104 | w.instances = newInstancesMap 105 | } 106 | 107 | func (w *sharedShardWriter) Close() { 108 | if !w.isClosed.CAS(false, true) { 109 | return 110 | } 111 | w.mw.Close() 112 | } 113 | 114 | func (w *sharedShardWriter) QueueSize() int { 115 | return w.mw.QueueSize() 116 | } 117 | 118 | func (w *sharedShardWriter) SetMessageTTLNanos(value int64) { 119 | w.mw.SetMessageTTLNanos(value) 120 | } 121 | 122 | type replicatedShardWriter struct { 123 | sync.RWMutex 124 | 125 | shard uint32 126 | numberOfShards uint32 127 | mPool messagePool 128 | opts Options 129 | logger log.Logger 130 | 131 | ackRouter ackRouter 132 | replicaID uint32 133 | messageTTLNanos int64 134 | messageWriters map[string]messageWriter 135 | isClosed bool 136 | m messageWriterMetrics 137 | } 138 | 139 | func newReplicatedShardWriter( 140 | shard, numberOfShards uint32, 141 | router ackRouter, 142 | mPool messagePool, 143 | opts Options, 144 | m messageWriterMetrics, 145 | ) shardWriter { 146 | return &replicatedShardWriter{ 147 | shard: shard, 148 | numberOfShards: numberOfShards, 149 | mPool: mPool, 150 | opts: opts, 151 | logger: opts.InstrumentOptions().Logger(), 152 | ackRouter: router, 153 | replicaID: 0, 154 | messageWriters: make(map[string]messageWriter), 155 | isClosed: false, 156 | m: m, 157 | } 158 | } 159 | 160 | func (w *replicatedShardWriter) Write(rm *producer.RefCountedMessage) { 161 | w.RLock() 162 | for _, mw := range w.messageWriters { 163 | mw.Write(rm) 164 | } 165 | w.RUnlock() 166 | } 167 | 168 | // This is not thread safe, must be called in one thread. 169 | func (w *replicatedShardWriter) UpdateInstances( 170 | instances []placement.Instance, 171 | cws map[string]consumerWriter, 172 | ) { 173 | // TODO: Schedule time after shardcutoff to clean up message writers that 174 | // are already cutoff. Otherwise it will wait until next placement change 175 | // to clean up. 176 | var ( 177 | newMessageWriters = make(map[string]messageWriter, len(instances)) 178 | toBeClosed []messageWriter 179 | toBeAdded = make(map[placement.Instance]consumerWriter, len(instances)) 180 | oldMessageWriters = w.messageWriters 181 | ) 182 | for _, instance := range instances { 183 | key := instance.Endpoint() 184 | if mw, ok := oldMessageWriters[key]; ok { 185 | newMessageWriters[key] = mw 186 | // Existing instance, try to update cutover cutoff times. 187 | w.updateCutoverCutoffNanos(mw, instance) 188 | continue 189 | } 190 | // This is a new instance. 191 | toBeAdded[instance] = cws[key] 192 | } 193 | for id, mw := range oldMessageWriters { 194 | if _, ok := newMessageWriters[id]; ok { 195 | // Still in the new placement. 196 | continue 197 | } 198 | // Keep the existing message writer and swap the consumer writer in it 199 | // with a new consumer writer in the placement update, so that the 200 | // messages buffered in the existing message writer can be tried on 201 | // the new consumer writer. 202 | if instance, cw, ok := anyKeyValueInMap(toBeAdded); ok { 203 | mw.AddConsumerWriter(cw) 204 | mw.RemoveConsumerWriter(id) 205 | w.updateCutoverCutoffNanos(mw, instance) 206 | newMessageWriters[instance.Endpoint()] = mw 207 | delete(toBeAdded, instance) 208 | continue 209 | } 210 | toBeClosed = append(toBeClosed, mw) 211 | } 212 | 213 | // If there are more instances for this shard, this happens when user 214 | // increased replication factor for the placement or just this shard. 215 | for instance, cw := range toBeAdded { 216 | replicatedShardID := uint64(w.replicaID*w.numberOfShards + w.shard) 217 | w.replicaID++ 218 | mw := newMessageWriter(replicatedShardID, w.mPool, w.opts, w.m) 219 | mw.AddConsumerWriter(cw) 220 | w.updateCutoverCutoffNanos(mw, instance) 221 | mw.Init() 222 | w.ackRouter.Register(replicatedShardID, mw) 223 | newMessageWriters[instance.Endpoint()] = mw 224 | } 225 | 226 | w.Lock() 227 | w.messageWriters = newMessageWriters 228 | w.setMessageTTLNanosWithLock(w.messageTTLNanos) 229 | w.Unlock() 230 | 231 | // If there are less instances for this shard, this happens when user 232 | // reduced replication factor for the placement or just this shard. 233 | for _, mw := range toBeClosed { 234 | mw := mw 235 | // This needs to be in done in a go routine as closing a message writer will 236 | // block until all messages consumed. 237 | go func() { 238 | mw.Close() 239 | w.ackRouter.Unregister(mw.ReplicatedShardID()) 240 | }() 241 | } 242 | } 243 | 244 | func (w *replicatedShardWriter) updateCutoverCutoffNanos( 245 | mw messageWriter, 246 | instance placement.Instance, 247 | ) { 248 | s, ok := instance.Shards().Shard(w.shard) 249 | if !ok { 250 | // Unexpected. 251 | w.logger.Errorf("could not find shard %d on instance %s", w.shard, instance.Endpoint()) 252 | return 253 | } 254 | mw.SetCutoffNanos(s.CutoffNanos()) 255 | mw.SetCutoverNanos(s.CutoverNanos()) 256 | } 257 | 258 | func (w *replicatedShardWriter) Close() { 259 | w.Lock() 260 | defer w.Unlock() 261 | 262 | if w.isClosed { 263 | return 264 | } 265 | w.isClosed = true 266 | for _, mw := range w.messageWriters { 267 | mw.Close() 268 | } 269 | } 270 | 271 | func (w *replicatedShardWriter) QueueSize() int { 272 | w.RLock() 273 | mws := w.messageWriters 274 | var l int 275 | for _, mw := range mws { 276 | l += mw.QueueSize() 277 | } 278 | w.RUnlock() 279 | return l 280 | } 281 | 282 | func (w *replicatedShardWriter) SetMessageTTLNanos(value int64) { 283 | w.Lock() 284 | w.messageTTLNanos = value 285 | w.setMessageTTLNanosWithLock(value) 286 | w.Unlock() 287 | } 288 | 289 | func (w *replicatedShardWriter) setMessageTTLNanosWithLock(value int64) { 290 | for _, mw := range w.messageWriters { 291 | mw.SetMessageTTLNanos(value) 292 | } 293 | } 294 | 295 | func anyKeyValueInMap( 296 | m map[placement.Instance]consumerWriter, 297 | ) (placement.Instance, consumerWriter, bool) { 298 | for key, value := range m { 299 | return key, value, true 300 | } 301 | return nil, nil, false 302 | } 303 | --------------------------------------------------------------------------------