├── .codecov.yml
├── .gitignore
├── .travis.yml
├── CODE_OF_CONDUCT.md
├── CONTRIBUTING.md
├── LICENSE
├── Makefile
├── README.md
├── admin.go
├── admin_test.go
├── check_license.sh
├── constants.go
├── cover.sh
├── data_accessor.go
├── data_accessor_test.go
├── glide.lock
├── glide.yaml
├── key_builder.go
├── model
├── constants.go
├── current_state.go
├── external_view.go
├── ideal_state.go
├── instance_config.go
├── live_instance.go
├── message.go
├── model_test.go
├── record.go
├── record_test.go
└── state_model_def.go
├── participant.go
├── participant_test.go
├── state_model.go
├── state_model_processor.go
├── state_model_test.go
├── test_participant.go
├── test_util.go
├── util
├── set.go
└── set_test.go
└── zk
├── client.go
├── client_test.go
├── embedded
├── start.sh
├── stop.sh
├── zk.cfg
└── zookeeper-3.4.9-fatjar.jar
├── fake_zk.go
├── fake_zk_conn.go
├── fake_zk_conn_test.go
└── test_util.go
/.codecov.yml:
--------------------------------------------------------------------------------
1 | coverage:
2 | range: 40..100
3 | round: down
4 | precision: 2
5 |
6 | status:
7 | project: # measuring the overall project coverage
8 | default: # context, you can create multiple ones with custom titles
9 | enabled: yes # must be yes|true to enable this status
10 | target: 40% # specify the target coverage for each commit status
11 | # option: "auto" (must increase from parent commit or pull request base)
12 | # option: "X%" a static target percentage to hit
13 | if_not_found: success # if parent is not found report status as success, error, or failure
14 | if_ci_failed: error # if ci fails report status as success, error, or failure
15 |
--------------------------------------------------------------------------------
/.gitignore:
--------------------------------------------------------------------------------
1 | # Ignore vendor created by glide
2 | vendor/
3 |
4 | # Ignore the dev zookeeper data
5 | zookeeper-data/
6 |
7 | # Ignore coverage file
8 | coverage.txt
9 |
10 | # IDE artifacts
11 | .idea/*
12 |
13 | # Ignore system files
14 | *.swp
15 | *.log
16 |
--------------------------------------------------------------------------------
/.travis.yml:
--------------------------------------------------------------------------------
1 | language: go
2 | go:
3 | - 1.9
4 | env:
5 | global:
6 | - TEST_TIMEOUT_SCALE=40
7 | cache:
8 | directories:
9 | - vendor
10 | install:
11 | - make dependencies
12 | script:
13 | - make lint
14 | - make test
15 | after_success:
16 | - make cover
17 | - bash <(curl -s https://codecov.io/bash)
18 |
--------------------------------------------------------------------------------
/CODE_OF_CONDUCT.md:
--------------------------------------------------------------------------------
1 | # Contributor Covenant Code of Conduct
2 |
3 | ## Our Pledge
4 |
5 | In the interest of fostering an open and welcoming environment, we as
6 | contributors and maintainers pledge to making participation in our project and
7 | our community a harassment-free experience for everyone, regardless of age,
8 | body size, disability, ethnicity, gender identity and expression, level of
9 | experience, nationality, personal appearance, race, religion, or sexual
10 | identity and orientation.
11 |
12 | ## Our Standards
13 |
14 | Examples of behavior that contributes to creating a positive environment
15 | include:
16 |
17 | * Using welcoming and inclusive language
18 | * Being respectful of differing viewpoints and experiences
19 | * Gracefully accepting constructive criticism
20 | * Focusing on what is best for the community
21 | * Showing empathy towards other community members
22 |
23 | Examples of unacceptable behavior by participants include:
24 |
25 | * The use of sexualized language or imagery and unwelcome sexual attention or
26 | advances
27 | * Trolling, insulting/derogatory comments, and personal or political attacks
28 | * Public or private harassment
29 | * Publishing others' private information, such as a physical or electronic
30 | address, without explicit permission
31 | * Other conduct which could reasonably be considered inappropriate in a
32 | professional setting
33 |
34 | ## Our Responsibilities
35 |
36 | Project maintainers are responsible for clarifying the standards of acceptable
37 | behavior and are expected to take appropriate and fair corrective action in
38 | response to any instances of unacceptable behavior.
39 |
40 | Project maintainers have the right and responsibility to remove, edit, or
41 | reject comments, commits, code, wiki edits, issues, and other contributions
42 | that are not aligned to this Code of Conduct, or to ban temporarily or
43 | permanently any contributor for other behaviors that they deem inappropriate,
44 | threatening, offensive, or harmful.
45 |
46 | ## Scope
47 |
48 | This Code of Conduct applies both within project spaces and in public spaces
49 | when an individual is representing the project or its community. Examples of
50 | representing a project or community include using an official project e-mail
51 | address, posting via an official social media account, or acting as an
52 | appointed representative at an online or offline event. Representation of a
53 | project may be further defined and clarified by project maintainers.
54 |
55 | ## Enforcement
56 |
57 | Instances of abusive, harassing, or otherwise unacceptable behavior may be
58 | reported by contacting the project team at oss-conduct@uber.com. The project
59 | team will review and investigate all complaints, and will respond in a way
60 | that it deems appropriate to the circumstances. The project team is obligated
61 | to maintain confidentiality with regard to the reporter of an incident.
62 | Further details of specific enforcement policies may be posted separately.
63 |
64 | Project maintainers who do not follow or enforce the Code of Conduct in good
65 | faith may face temporary or permanent repercussions as determined by other
66 | members of the project's leadership.
67 |
68 | ## Attribution
69 |
70 | This Code of Conduct is adapted from the [Contributor Covenant][homepage],
71 | version 1.4, available at
72 | [http://contributor-covenant.org/version/1/4][version].
73 |
74 | [homepage]: http://contributor-covenant.org
75 | [version]: http://contributor-covenant.org/version/1/4/
76 |
--------------------------------------------------------------------------------
/CONTRIBUTING.md:
--------------------------------------------------------------------------------
1 | # Contributing
2 |
3 | We'd love your help making the project better.
4 |
5 | In your issue, pull request, and any other
6 | communications, please remember to treat your fellow contributors with
7 | respect! We take our [code of conduct](CODE_OF_CONDUCT.md) seriously.
8 |
9 | ## Setup
10 |
11 | [Fork][fork], and then clone the repository:
12 |
13 | ```
14 | mkdir -p $GOPATH/src/github.com/uber-go
15 | cd $GOPATH/src/github.com/uber-go
16 | git clone git@github.com:your_github_username/go-helix.git
17 | cd go-helix
18 | git remote add upstream https://github.com/uber-go/go-helix.git
19 | git fetch upstream
20 | ```
21 |
22 | Install go-helix's dependencies:
23 |
24 | ```
25 | make dependencies
26 | ```
27 |
28 | Make sure that the tests and the linters pass:
29 |
30 | ```
31 | make test
32 | make lint
33 | ```
34 |
35 | If you're not using the minor version of Go specified in the Makefile's
36 | `LINTABLE_MINOR_VERSIONS` variable, `make lint` doesn't do anything. This is
37 | fine, but it means that you'll only discover lint failures after you open your
38 | pull request.
39 |
40 | ## Making Changes
41 |
42 | Start by creating a new branch for your changes:
43 |
44 | ```
45 | cd $GOPATH/src/github.com/uber-go/go-helix
46 | git checkout master
47 | git fetch upstream
48 | git rebase upstream/master
49 | git checkout -b cool_new_feature
50 | ```
51 |
52 | Make your changes, then ensure that `make lint` and `make test` still pass. If
53 | you're satisfied with your changes, push them to your fork.
54 |
55 | ```
56 | git push origin cool_new_feature
57 | ```
58 |
59 | Then use the GitHub UI to open a pull request.
60 |
61 | At this point, you're waiting on us to review your changes. We *try* to respond
62 | to issues and pull requests within a few business days, and we may suggest some
63 | improvements or alternatives. Once your changes are approved, one of the
64 | project maintainers will merge them.
65 |
66 | We're much more likely to approve your changes if you:
67 |
68 | * Add tests for new functionality.
69 | * Write a [good commit message][commit-message].
70 | * Maintain backward compatibility.
71 |
72 | [fork]: https://github.com/uber-go/go-helix/fork
73 | [open-issue]: https://github.com/uber-go/go-helix/issues/new
74 | [commit-message]: http://tbaggery.com/2008/04/19/a-note-about-git-commit-messages.html
75 |
--------------------------------------------------------------------------------
/LICENSE:
--------------------------------------------------------------------------------
1 | go-helix
2 | Copyright (c) 2017 Uber Technologies, Inc.
3 |
4 | Permission is hereby granted, free of charge, to any person obtaining a copy
5 | of this software and associated documentation files (the "Software"), to deal
6 | in the Software without restriction, including without limitation the rights
7 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
8 | copies of the Software, and to permit persons to whom the Software is
9 | furnished to do so, subject to the following conditions:
10 |
11 | The above copyright notice and this permission notice shall be included in all
12 | copies or substantial portions of the Software.
13 |
14 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
15 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
16 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
17 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
18 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
19 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
20 | SOFTWARE.
21 |
22 | -------
23 |
24 | go-helix is based, at least in part, on the following:
25 |
26 | Apache Helix
27 |
28 | Copyright (c) 2017 Apache Software Foundation
29 |
30 | Licensed under the Apache License, Version 2.0 (the "License");
31 | you may not use this file except in compliance with the License.
32 | You may obtain a copy of the License at
33 |
34 | http://www.apache.org/licenses/LICENSE-2.0
35 |
36 | Unless required by applicable law or agreed to in writing, software
37 | distributed under the License is distributed on an "AS IS" BASIS,
38 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
39 | See the License for the specific language governing permissions and
40 | limitations under the License.
41 |
42 | -------
43 |
44 | go-helix includes or depends on the following:
45 |
46 | Zap
47 | Copyright (c) 2016-2017 Uber Technologies, Inc.
48 |
49 | Permission is hereby granted, free of charge, to any person obtaining a copy
50 | of this software and associated documentation files (the "Software"), to deal
51 | in the Software without restriction, including without limitation the rights
52 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
53 | copies of the Software, and to permit persons to whom the Software is
54 | furnished to do so, subject to the following conditions:
55 |
56 | The above copyright notice and this permission notice shall be included in all
57 | copies or substantial portions of the Software.
58 |
59 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
60 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
61 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
62 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
63 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
64 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
65 | SOFTWARE.
66 |
67 | Tally
68 | Copyright (c) 2016 Uber Technologies, Inc.
69 |
70 | Permission is hereby granted, free of charge, to any person obtaining a copy
71 | of this software and associated documentation files (the "Software"), to deal
72 | in the Software without restriction, including without limitation the rights
73 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
74 | copies of the Software, and to permit persons to whom the Software is
75 | furnished to do so, subject to the following conditions:
76 |
77 | The above copyright notice and this permission notice shall be included in all
78 | copies or substantial portions of the Software.
79 |
80 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
81 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
82 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
83 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
84 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
85 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
86 | SOFTWARE.
87 |
--------------------------------------------------------------------------------
/Makefile:
--------------------------------------------------------------------------------
1 | PKGS ?= $(shell glide novendor)
2 | # Many Go tools take file globs or directories as arguments instead of packages.
3 | PKG_FILES ?= *.go model util zk
4 |
5 | # The linting tools evolve with each Go version, so run them only on the latest
6 | # stable release.
7 | GO_VERSION := $(shell go version | cut -d " " -f 3)
8 | GO_MINOR_VERSION := $(word 2,$(subst ., ,$(GO_VERSION)))
9 | LINTABLE_MINOR_VERSIONS := 9
10 | ifneq ($(filter $(LINTABLE_MINOR_VERSIONS),$(GO_MINOR_VERSION)),)
11 | SHOULD_LINT := true
12 | endif
13 |
14 | .PHONY: all
15 | all: lint test
16 |
17 | .PHONY: lint
18 | lint:
19 | ifdef SHOULD_LINT
20 | @rm -rf lint.log
21 | @echo "Checking formatting..."
22 | @gofmt -d -s $(PKG_FILES) 2>&1 | tee lint.log
23 | @echo "Installing test dependencies for vet..."
24 | @go test -i $(PKGS)
25 | @echo "Checking vet..."
26 | @$(foreach dir,$(PKG_FILES),go tool vet -printf=false $(dir) 2>&1 | tee -a lint.log;)
27 | @echo "Checking lint..."
28 | @$(foreach dir,$(PKGS),golint $(dir) 2>&1 | tee -a lint.log;)
29 | @echo "Checking for license headers..."
30 | @./check_license.sh | tee -a lint.log
31 | @[ ! -s lint.log ]
32 | else
33 | @echo "Skipping linters on" $(GO_VERSION)
34 | endif
35 |
36 | .PHONY: test
37 | test:
38 | go test -p 1 -race $(PKGS)
39 |
40 | .PHONY: cover
41 | cover:
42 | ./cover.sh $(PKGS)
43 |
44 | .PHONY: dependencies
45 | dependencies:
46 | @echo "Installing Glide and locked dependencies..."
47 | glide --version || go get -u -f github.com/Masterminds/glide
48 | glide install
49 | ifdef SHOULD_LINT
50 | @echo "Installing golint..."
51 | go install ./vendor/github.com/golang/lint/golint
52 | else
53 | @echo "Not installing golint, since we don't expect to lint on" $(GO_VERSION)
54 | endif
55 |
56 | # utility to find project root
57 | export APP_ROOT = $(CURDIR)
58 |
--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
1 | # go-helix [![GoDoc][doc-img]][doc] [![Build Status][ci-img]][ci] [![Coverage Status][cov-img]][cov]
2 |
3 | A Go implementation of [Apache Helix](https://helix.apache.org).
4 |
5 | Currently the participant part only, compatible with the Apache Helix Java controller.
6 |
7 | ## Installation
8 |
9 | `go get -u github.com/uber-go/go-helix`
10 |
11 | ## Quick Start
12 |
13 | ### Init participant
14 |
15 | ```go
16 | participant, fatalErrChan := NewParticipant(
17 | zap.NewNop(),
18 | tally.NoopScope,
19 | "localhost:2181", // Zookeeper connect string
20 | "test_app", // application identifier
21 | "test_cluster", // helix cluster name
22 | "test_resource", // helix resource name
23 | "localhost", // participant host name
24 | 123, // participant port
25 | )
26 |
27 | processor := NewStateModelProcessor()
28 | processor.AddTransition("OFFLINE", "ONLINE", func(m *model.Message) (err error) {
29 | partitionNum, err := p.getPartitionNum(m)
30 | if err != nil {
31 | return err
32 | }
33 | // add logic here to save the owned partition and/or perform any actions for going online
34 | })
35 | processor.AddTransition("ONLINE", "OFFLINE", func(m *model.Message) (err error) {
36 | partitionNum, err := p.getPartitionNum(m)
37 | if err != nil {
38 | return err
39 | }
40 | // add logic here to remove the owned partition and/or perform any actions for going offline
41 | })
42 |
43 | participant.RegisterStateModel("OnlineOffline", processor)
44 |
45 | err := participant.Connect() // initialization is complete if err is nil
46 |
47 | // listen to the fatal error channel and handle the error by
48 | // 1. recreate participant from scratch and connect, or
49 | // 2. quit the program and restart
50 | fatalErr := <- fatalErrChan
51 | ```
52 |
53 | ### Use participant
54 |
55 | Use the saved partitions to see if the partition should be handled by the participant.
56 |
57 | ### Disconnect participant
58 |
59 | ```go
60 | participant.Disconnect()
61 | ```
62 |
63 | ## Development Status: Beta
64 |
65 | The APIs are functional. We do not expect, but there's no guarantee that no breaking changes will be made.
66 |
67 | ## Contributing
68 |
69 | We encourage and support an active, healthy community of contributors —
70 | including you! Details are in the [contribution guide](CONTRIBUTING.md) and
71 | the [code of conduct](CODE_OF_CONDUCT.md). The go-helix maintainers keep an eye on
72 | issues and pull requests, but you can also report any negative conduct to
73 | oss-conduct@uber.com. That email list is a private, safe space; even the go-helix
74 | maintainers don't have access, so don't hesitate to hold us to a high standard.
75 |
76 |
77 |
78 | Released under the [MIT License](LICENSE).
79 |
80 | [doc-img]: https://godoc.org/github.com/uber-go/go-helix?status.svg
81 | [doc]: https://godoc.org/github.com/uber-go/go-helix
82 | [ci-img]: https://travis-ci.org/uber-go/go-helix.svg?branch=master
83 | [ci]: https://travis-ci.org/uber-go/go-helix
84 | [cov-img]: https://codecov.io/gh/uber-go/go-helix/branch/master/graph/badge.svg
85 | [cov]: https://codecov.io/gh/uber-go/go-helix
86 |
--------------------------------------------------------------------------------
/admin.go:
--------------------------------------------------------------------------------
1 | // Copyright (c) 2017 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 helix
22 |
23 | import (
24 | "bytes"
25 | "fmt"
26 | "strconv"
27 | "strings"
28 |
29 | "github.com/pkg/errors"
30 | "github.com/uber-go/go-helix/model"
31 | "github.com/uber-go/go-helix/zk"
32 | "github.com/uber-go/tally"
33 | "go.uber.org/zap"
34 | )
35 |
36 | var (
37 | // ErrClusterNotSetup means the helix data structure in zookeeper /{CLUSTER_NAME}
38 | // is not correct or does not exist
39 | ErrClusterNotSetup = errors.New("cluster not setup")
40 |
41 | // ErrNodeAlreadyExists the zookeeper node exists when it is not expected to
42 | ErrNodeAlreadyExists = errors.New("node already exists in cluster")
43 |
44 | // ErrNodeNotExist the zookeeper node does not exist when it is expected to
45 | ErrNodeNotExist = errors.New("node does not exist in config for cluster")
46 |
47 | // ErrInstanceNotExist the instance of a cluster does not exist when it is expected to
48 | ErrInstanceNotExist = errors.New("node does not exist in instances for cluster")
49 |
50 | // ErrStateModelDefNotExist the state model definition is expected to exist in zookeeper
51 | ErrStateModelDefNotExist = errors.New("state model not exist in cluster")
52 |
53 | // ErrResourceExists the resource already exists in cluster and cannot be added again
54 | ErrResourceExists = errors.New("resource already exists in cluster")
55 |
56 | // ErrResourceNotExists the resource does not exists and cannot be removed
57 | ErrResourceNotExists = errors.New("resource not exists in cluster")
58 | )
59 |
60 | var (
61 | _helixDefaultNodes = map[string]string{
62 | "LeaderStandby": `
63 | {
64 | "id" : "LeaderStandby",
65 | "mapFields" : {
66 | "DROPPED.meta" : {
67 | "count" : "-1"
68 | },
69 | "LEADER.meta" : {
70 | "count" : "1"
71 | },
72 | "LEADER.next" : {
73 | "DROPPED" : "STANDBY",
74 | "STANDBY" : "STANDBY",
75 | "OFFLINE" : "STANDBY"
76 | },
77 | "OFFLINE.meta" : {
78 | "count" : "-1"
79 | },
80 | "OFFLINE.next" : {
81 | "DROPPED" : "DROPPED",
82 | "STANDBY" : "STANDBY",
83 | "LEADER" : "STANDBY"
84 | },
85 | "STANDBY.meta" : {
86 | "count" : "R"
87 | },
88 | "STANDBY.next" : {
89 | "DROPPED" : "OFFLINE",
90 | "OFFLINE" : "OFFLINE",
91 | "LEADER" : "LEADER"
92 | }
93 | },
94 | "listFields" : {
95 | "STATE_PRIORITY_LIST" : [ "LEADER", "STANDBY", "OFFLINE", "DROPPED" ],
96 | "STATE_TRANSITION_PRIORITYLIST" : [ "LEADER-STANDBY", "STANDBY-LEADER", "OFFLINE-STANDBY", "STANDBY-OFFLINE", "OFFLINE-DROPPED" ]
97 | },
98 | "simpleFields" : {
99 | "INITIAL_STATE" : "OFFLINE"
100 | }
101 | }
102 | `,
103 | "MasterSlave": `
104 | {
105 | "id" : "MasterSlave",
106 | "mapFields" : {
107 | "DROPPED.meta" : {
108 | "count" : "-1"
109 | },
110 | "ERROR.meta" : {
111 | "count" : "-1"
112 | },
113 | "ERROR.next" : {
114 | "DROPPED" : "DROPPED",
115 | "OFFLINE" : "OFFLINE"
116 | },
117 | "MASTER.meta" : {
118 | "count" : "1"
119 | },
120 | "MASTER.next" : {
121 | "SLAVE" : "SLAVE",
122 | "DROPPED" : "SLAVE",
123 | "OFFLINE" : "SLAVE"
124 | },
125 | "OFFLINE.meta" : {
126 | "count" : "-1"
127 | },
128 | "OFFLINE.next" : {
129 | "SLAVE" : "SLAVE",
130 | "DROPPED" : "DROPPED",
131 | "MASTER" : "SLAVE"
132 | },
133 | "SLAVE.meta" : {
134 | "count" : "R"
135 | },
136 | "SLAVE.next" : {
137 | "DROPPED" : "OFFLINE",
138 | "OFFLINE" : "OFFLINE",
139 | "MASTER" : "MASTER"
140 | }
141 | },
142 | "listFields" : {
143 | "STATE_PRIORITY_LIST" : [ "MASTER", "SLAVE", "OFFLINE", "DROPPED", "ERROR" ],
144 | "STATE_TRANSITION_PRIORITYLIST" : [ "MASTER-SLAVE", "SLAVE-MASTER", "OFFLINE-SLAVE", "SLAVE-OFFLINE", "OFFLINE-DROPPED" ]
145 | },
146 | "simpleFields" : {
147 | "INITIAL_STATE" : "OFFLINE"
148 | }
149 | }
150 | `,
151 | "OnlineOffline": `
152 | {
153 | "id" : "OnlineOffline",
154 | "mapFields" : {
155 | "DROPPED.meta" : {
156 | "count" : "-1"
157 | },
158 | "OFFLINE.meta" : {
159 | "count" : "-1"
160 | },
161 | "OFFLINE.next" : {
162 | "DROPPED" : "DROPPED",
163 | "ONLINE" : "ONLINE"
164 | },
165 | "ONLINE.meta" : {
166 | "count" : "R"
167 | },
168 | "ONLINE.next" : {
169 | "DROPPED" : "OFFLINE",
170 | "OFFLINE" : "OFFLINE"
171 | }
172 | },
173 | "listFields" : {
174 | "STATE_PRIORITY_LIST" : [ "ONLINE", "OFFLINE", "DROPPED" ],
175 | "STATE_TRANSITION_PRIORITYLIST" : [ "OFFLINE-ONLINE", "ONLINE-OFFLINE", "OFFLINE-DROPPED" ]
176 | },
177 | "simpleFields" : {
178 | "INITIAL_STATE" : "OFFLINE"
179 | }
180 | }
181 | `,
182 | "STORAGE_DEFAULT_SM_SCHEMATA": `
183 | {
184 | "id" : "STORAGE_DEFAULT_SM_SCHEMATA",
185 | "mapFields" : {
186 | "DROPPED.meta" : {
187 | "count" : "-1"
188 | },
189 | "ERROR.meta" : {
190 | "count" : "-1"
191 | },
192 | "ERROR.next" : {
193 | "DROPPED" : "DROPPED",
194 | "OFFLINE" : "OFFLINE"
195 | },
196 | "MASTER.meta" : {
197 | "count" : "N"
198 | },
199 | "MASTER.next" : {
200 | "DROPPED" : "OFFLINE",
201 | "OFFLINE" : "OFFLINE"
202 | },
203 | "OFFLINE.meta" : {
204 | "count" : "-1"
205 | },
206 | "OFFLINE.next" : {
207 | "DROPPED" : "DROPPED",
208 | "MASTER" : "MASTER"
209 | }
210 | },
211 | "listFields" : {
212 | "STATE_PRIORITY_LIST" : [ "MASTER", "OFFLINE", "DROPPED", "ERROR" ],
213 | "STATE_TRANSITION_PRIORITYLIST" : [ "MASTER-OFFLINE", "OFFLINE-MASTER" ]
214 | },
215 | "simpleFields" : {
216 | "INITIAL_STATE" : "OFFLINE"
217 | }
218 | }
219 | `,
220 | "SchedulerTaskQueue": `
221 | {
222 | "id" : "SchedulerTaskQueue",
223 | "mapFields" : {
224 | "COMPLETED.meta" : {
225 | "count" : "1"
226 | },
227 | "COMPLETED.next" : {
228 | "DROPPED" : "DROPPED",
229 | "COMPLETED" : "COMPLETED"
230 | },
231 | "DROPPED.meta" : {
232 | "count" : "-1"
233 | },
234 | "DROPPED.next" : {
235 | "DROPPED" : "DROPPED"
236 | },
237 | "OFFLINE.meta" : {
238 | "count" : "-1"
239 | },
240 | "OFFLINE.next" : {
241 | "DROPPED" : "DROPPED",
242 | "OFFLINE" : "OFFLINE",
243 | "COMPLETED" : "COMPLETED"
244 | }
245 | },
246 | "listFields" : {
247 | "STATE_PRIORITY_LIST" : [ "COMPLETED", "OFFLINE", "DROPPED" ],
248 | "STATE_TRANSITION_PRIORITYLIST" : [ "OFFLINE-COMPLETED", "OFFLINE-DROPPED", "COMPLETED-DROPPED" ]
249 | },
250 | "simpleFields" : {
251 | "INITIAL_STATE" : "OFFLINE"
252 | }
253 | }
254 | `,
255 | "Task": `
256 | {
257 | "id" : "Task",
258 | "mapFields" : {
259 | "COMPLETED.meta" : {
260 | "count" : "-1"
261 | },
262 | "COMPLETED.next" : {
263 | "STOPPED" : "INIT",
264 | "DROPPED" : "DROPPED",
265 | "RUNNING" : "INIT",
266 | "INIT" : "INIT",
267 | "COMPLETED" : "COMPLETED",
268 | "TASK_ERROR" : "INIT",
269 | "TIMED_OUT" : "INIT"
270 | },
271 | "DROPPED.meta" : {
272 | "count" : "-1"
273 | },
274 | "DROPPED.next" : {
275 | "DROPPED" : "DROPPED"
276 | },
277 | "INIT.meta" : {
278 | "count" : "-1"
279 | },
280 | "INIT.next" : {
281 | "STOPPED" : "RUNNING",
282 | "DROPPED" : "DROPPED",
283 | "RUNNING" : "RUNNING",
284 | "INIT" : "INIT",
285 | "COMPLETED" : "RUNNING",
286 | "TASK_ERROR" : "RUNNING",
287 | "TIMED_OUT" : "RUNNING"
288 | },
289 | "RUNNING.meta" : {
290 | "count" : "-1"
291 | },
292 | "RUNNING.next" : {
293 | "STOPPED" : "STOPPED",
294 | "DROPPED" : "DROPPED",
295 | "RUNNING" : "RUNNING",
296 | "INIT" : "INIT",
297 | "COMPLETED" : "COMPLETED",
298 | "TASK_ERROR" : "TASK_ERROR",
299 | "TIMED_OUT" : "TIMED_OUT"
300 | },
301 | "STOPPED.meta" : {
302 | "count" : "-1"
303 | },
304 | "STOPPED.next" : {
305 | "STOPPED" : "STOPPED",
306 | "DROPPED" : "DROPPED",
307 | "RUNNING" : "RUNNING",
308 | "INIT" : "INIT",
309 | "COMPLETED" : "RUNNING",
310 | "TASK_ERROR" : "RUNNING",
311 | "TIMED_OUT" : "RUNNING"
312 | },
313 | "TASK_ERROR.meta" : {
314 | "count" : "-1"
315 | },
316 | "TASK_ERROR.next" : {
317 | "STOPPED" : "INIT",
318 | "DROPPED" : "DROPPED",
319 | "RUNNING" : "INIT",
320 | "INIT" : "INIT",
321 | "COMPLETED" : "INIT",
322 | "TIMED_OUT" : "INIT",
323 | "TASK_ERROR" : "TASK_ERROR"
324 | },
325 | "TIMED_OUT.meta" : {
326 | "count" : "-1"
327 | },
328 | "TIMED_OUT.next" : {
329 | "STOPPED" : "INIT",
330 | "DROPPED" : "DROPPED",
331 | "RUNNING" : "INIT",
332 | "INIT" : "INIT",
333 | "COMPLETED" : "INIT",
334 | "TASK_ERROR" : "INIT",
335 | "TIMED_OUT" : "TIMED_OUT"
336 | }
337 | },
338 | "listFields" : {
339 | "STATE_PRIORITY_LIST" : [ "INIT", "RUNNING", "STOPPED", "COMPLETED", "TIMED_OUT", "TASK_ERROR", "DROPPED" ],
340 | "STATE_TRANSITION_PRIORITYLIST" : [ "INIT-RUNNING", "RUNNING-STOPPED", "RUNNING-COMPLETED", "RUNNING-TIMED_OUT", "RUNNING-TASK_ERROR", "STOPPED-RUNNING", "INIT-DROPPED", "RUNNING-DROPPED", "COMPLETED-DROPPED", "STOPPED-DROPPED", "TIMED_OUT-DROPPED", "TASK_ERROR-DROPPED", "RUNNING-INIT", "COMPLETED-INIT", "STOPPED-INIT", "TIMED_OUT-INIT", "TASK_ERROR-INIT" ]
341 | },
342 | "simpleFields" : {
343 | "INITIAL_STATE" : "INIT"
344 | }
345 | }`,
346 | }
347 | )
348 |
349 | // Admin handles the administration task for the Helix cluster. Many of the operations
350 | // are mirroring the implementions documented at
351 | // http://helix.apache.org/0.7.0-incubating-docs/Quickstart.html
352 | type Admin struct {
353 | zkClient *zk.Client
354 | }
355 |
356 | // NewAdmin instantiates Admin
357 | func NewAdmin(zkConnectString string) (*Admin, error) {
358 | zkClient := zk.NewClient(
359 | zap.NewNop(), tally.NoopScope, zk.WithZkSvr(zkConnectString), zk.WithSessionTimeout(zk.DefaultSessionTimeout))
360 | err := zkClient.Connect()
361 | if err != nil {
362 | return nil, err
363 | }
364 |
365 | return &Admin{zkClient: zkClient}, nil
366 | }
367 |
368 | // AddCluster add a cluster to Helix. As a result, a znode will be created in zookeeper
369 | // root named after the cluster name, and corresponding data structures are populated
370 | // under this znode.
371 | // The cluster would be dropped and recreated if recreateIfExists is true
372 | func (adm Admin) AddCluster(cluster string, recreateIfExists bool) bool {
373 | kb := &KeyBuilder{cluster}
374 | // c = "/"
375 | c := kb.cluster()
376 |
377 | // check if cluster already exists
378 | exists, _, err := adm.zkClient.Exists(c)
379 | if err != nil || (exists && !recreateIfExists) {
380 | return false
381 | }
382 |
383 | if recreateIfExists {
384 | if err := adm.zkClient.DeleteTree(c); err != nil {
385 | return false
386 | }
387 | }
388 |
389 | adm.zkClient.CreateEmptyNode(c)
390 |
391 | // PROPERTYSTORE is an empty node
392 | propertyStore := fmt.Sprintf("/%s/PROPERTYSTORE", cluster)
393 | adm.zkClient.CreateEmptyNode(propertyStore)
394 |
395 | // STATEMODELDEFS has 6 children
396 | stateModelDefs := fmt.Sprintf("/%s/STATEMODELDEFS", cluster)
397 | adm.zkClient.CreateEmptyNode(stateModelDefs)
398 | adm.zkClient.CreateDataWithPath(
399 | stateModelDefs+"/LeaderStandby", []byte(_helixDefaultNodes["LeaderStandby"]))
400 | adm.zkClient.CreateDataWithPath(
401 | stateModelDefs+"/MasterSlave", []byte(_helixDefaultNodes["MasterSlave"]))
402 | adm.zkClient.CreateDataWithPath(
403 | stateModelDefs+"/OnlineOffline", []byte(_helixDefaultNodes[StateModelNameOnlineOffline]))
404 | adm.zkClient.CreateDataWithPath(
405 | stateModelDefs+"/STORAGE_DEFAULT_SM_SCHEMATA",
406 | []byte(_helixDefaultNodes["STORAGE_DEFAULT_SM_SCHEMATA"]))
407 | adm.zkClient.CreateDataWithPath(
408 | stateModelDefs+"/SchedulerTaskQueue", []byte(_helixDefaultNodes["SchedulerTaskQueue"]))
409 | adm.zkClient.CreateDataWithPath(
410 | stateModelDefs+"/Task", []byte(_helixDefaultNodes["Task"]))
411 |
412 | // INSTANCES is initailly an empty node
413 | instances := fmt.Sprintf("/%s/INSTANCES", cluster)
414 | adm.zkClient.CreateEmptyNode(instances)
415 |
416 | // CONFIGS has 3 children: CLUSTER, RESOURCE, PARTICIPANT
417 | configs := fmt.Sprintf("/%s/CONFIGS", cluster)
418 | adm.zkClient.CreateEmptyNode(configs)
419 | adm.zkClient.CreateEmptyNode(configs + "/PARTICIPANT")
420 | adm.zkClient.CreateEmptyNode(configs + "/RESOURCE")
421 | adm.zkClient.CreateEmptyNode(configs + "/CLUSTER")
422 |
423 | clusterNode := model.NewMsg(cluster)
424 | accessor := newDataAccessor(adm.zkClient, kb)
425 | accessor.createMsg(configs+"/CLUSTER/"+cluster, clusterNode)
426 |
427 | // empty ideal states
428 | idealStates := fmt.Sprintf("/%s/IDEALSTATES", cluster)
429 | adm.zkClient.CreateEmptyNode(idealStates)
430 |
431 | // empty external view
432 | externalView := fmt.Sprintf("/%s/EXTERNALVIEW", cluster)
433 | adm.zkClient.CreateEmptyNode(externalView)
434 |
435 | // empty live instances
436 | liveInstances := fmt.Sprintf("/%s/LIVEINSTANCES", cluster)
437 | adm.zkClient.CreateEmptyNode(liveInstances)
438 |
439 | // CONTROLLER has four childrens: [ERRORS, HISTORY, MESSAGES, STATUSUPDATES]
440 | controller := fmt.Sprintf("/%s/CONTROLLER", cluster)
441 | adm.zkClient.CreateEmptyNode(controller)
442 | adm.zkClient.CreateEmptyNode(controller + "/ERRORS")
443 | adm.zkClient.CreateEmptyNode(controller + "/HISTORY")
444 | adm.zkClient.CreateEmptyNode(controller + "/MESSAGES")
445 | adm.zkClient.CreateEmptyNode(controller + "/STATUSUPDATES")
446 |
447 | return true
448 | }
449 |
450 | // SetConfig set the configuration values for the cluster, defined by the config scope
451 | func (adm Admin) SetConfig(cluster string, scope string, properties map[string]string) error {
452 | switch strings.ToUpper(scope) {
453 | case "CLUSTER":
454 | if allow, ok := properties[_allowParticipantAutoJoinKey]; ok {
455 | builder := KeyBuilder{cluster}
456 | path := builder.clusterConfig()
457 |
458 | if strings.ToLower(allow) == "true" {
459 | adm.zkClient.UpdateSimpleField(path, _allowParticipantAutoJoinKey, "true")
460 | }
461 | }
462 | case "CONSTRAINT":
463 | case "PARTICIPANT":
464 | case "PARTITION":
465 | case "RESOURCE":
466 | }
467 |
468 | return nil
469 | }
470 |
471 | // GetConfig obtains the configuration value of a property, defined by a config scope
472 | func (adm Admin) GetConfig(
473 | cluster string, scope string, builder []string) (map[string]interface{}, error) {
474 | result := make(map[string]interface{})
475 |
476 | switch scope {
477 | case "CLUSTER":
478 | kb := KeyBuilder{cluster}
479 | path := kb.clusterConfig()
480 |
481 | for _, k := range builder {
482 | var err error
483 | val, err := adm.zkClient.GetSimpleFieldValueByKey(path, k)
484 | if err != nil {
485 | return nil, err
486 | }
487 | result[k] = val
488 | }
489 | case "CONSTRAINT":
490 | case "PARTICIPANT":
491 | case "PARTITION":
492 | case "RESOURCE":
493 | }
494 |
495 | return result, nil
496 | }
497 |
498 | // DropCluster removes a helix cluster from zookeeper. This will remove the
499 | // znode named after the cluster name from the zookeeper root.
500 | func (adm Admin) DropCluster(cluster string) error {
501 | kb := KeyBuilder{cluster}
502 | c := kb.cluster()
503 |
504 | return adm.zkClient.DeleteTree(c)
505 | }
506 |
507 | // AddNode is the internal implementation corresponding to command
508 | // ./helix-admin.sh --zkSvr --addNode
509 | // node is in the form of host_port
510 | func (adm Admin) AddNode(cluster string, node string) error {
511 | if ok, err := adm.isClusterSetup(cluster); ok == false || err != nil {
512 | return ErrClusterNotSetup
513 | }
514 |
515 | // check if node already exists under //CONFIGS/PARTICIPANT/
516 | builder := &KeyBuilder{cluster}
517 | path := builder.participantConfig(node)
518 | exists, _, err := adm.zkClient.Exists(path)
519 | if err != nil {
520 | return err
521 | }
522 | if exists {
523 | return ErrNodeAlreadyExists
524 | }
525 |
526 | // create new node for the participant
527 | parts := strings.Split(node, "_")
528 | n := model.NewMsg(node)
529 | n.SetSimpleField("HELIX_HOST", parts[0])
530 | n.SetSimpleField("HELIX_PORT", parts[1])
531 |
532 | accessor := newDataAccessor(adm.zkClient, builder)
533 | accessor.createMsg(path, n)
534 | adm.zkClient.CreateEmptyNode(builder.instance(node))
535 | adm.zkClient.CreateEmptyNode(builder.participantMessages(node))
536 | adm.zkClient.CreateEmptyNode(builder.currentStates(node))
537 | adm.zkClient.CreateEmptyNode(builder.errorsR(node))
538 | adm.zkClient.CreateEmptyNode(builder.statusUpdates(node))
539 |
540 | return nil
541 | }
542 |
543 | // DropNode removes a node from a cluster. The corresponding znodes
544 | // in zookeeper will be removed.
545 | func (adm Admin) DropNode(cluster string, node string) error {
546 | // check if node already exists under //CONFIGS/PARTICIPANT/
547 | builder := KeyBuilder{cluster}
548 | if exists, _, err := adm.zkClient.Exists(builder.participantConfig(node)); !exists || err != nil {
549 | return ErrNodeNotExist
550 | }
551 |
552 | // check if node exist under instance: //INSTANCES/
553 | if exists, _, err := adm.zkClient.Exists(builder.instance(node)); !exists || err != nil {
554 | return ErrInstanceNotExist
555 | }
556 |
557 | // delete //CONFIGS/PARTICIPANT/
558 | adm.zkClient.DeleteTree(builder.participantConfig(node))
559 |
560 | // delete //INSTANCES/
561 | adm.zkClient.DeleteTree(builder.instance(node))
562 |
563 | return nil
564 | }
565 |
566 | // AddResource implements the helix-admin.sh --addResource
567 | // ./helix-admin.sh --zkSvr localhost:2199 --addResource MYCLUSTER myDB 6 MasterSlave
568 | func (adm Admin) AddResource(
569 | cluster string, resource string, partitions int, stateModel string) error {
570 | if ok, err := adm.isClusterSetup(cluster); !ok || err != nil {
571 | return ErrClusterNotSetup
572 | }
573 |
574 | builder := &KeyBuilder{cluster}
575 |
576 | // make sure the state model def exists
577 | exists, _, err := adm.zkClient.Exists(builder.stateModelDef(stateModel))
578 | if err != nil {
579 | return errors.Wrap(err, "state model doesnt't exist "+stateModel)
580 | }
581 | if !exists {
582 | return ErrStateModelDefNotExist
583 | }
584 |
585 | // make sure the path for the ideal state does not exit
586 | isPath := builder.idealStates() + "/" + resource
587 | if exists, _, err := adm.zkClient.Exists(isPath); exists || err != nil {
588 | if exists {
589 | return ErrResourceExists
590 | }
591 | return err
592 | }
593 |
594 | // create the idealstate for the resource
595 | // is := NewIdealState(resource)
596 | // is.SetNumPartitions(partitions)
597 | // is.SetReplicas(0)
598 | // is.SetRebalanceMode("SEMI_AUTO")
599 | // is.SetStateModelDefRef(stateModel)
600 | // // save the ideal state in zookeeper
601 | // is.Save(conn, cluster)
602 |
603 | is := model.NewMsg(resource)
604 | is.SetSimpleField("NUM_PARTITIONS", strconv.Itoa(partitions))
605 | is.SetSimpleField("REPLICAS", strconv.Itoa(0))
606 | is.SetSimpleField("REBALANCE_MODE", strings.ToUpper("SEMI_AUTO"))
607 | is.SetStateModelDef(stateModel)
608 |
609 | accessor := newDataAccessor(adm.zkClient, builder)
610 | accessor.createMsg(isPath, is)
611 |
612 | return nil
613 | }
614 |
615 | // DropResource removes the specified resource from the cluster.
616 | func (adm Admin) DropResource(cluster string, resource string) error {
617 | // make sure the cluster is already setup
618 | if ok, err := adm.isClusterSetup(cluster); !ok || err != nil {
619 | return ErrClusterNotSetup
620 | }
621 |
622 | builder := KeyBuilder{cluster}
623 |
624 | // make sure the path for the ideal state does not exit
625 | adm.zkClient.DeleteTree(builder.idealStates() + "/" + resource)
626 | adm.zkClient.DeleteTree(builder.resourceConfig(resource))
627 |
628 | return nil
629 | }
630 |
631 | // EnableResource enables the specified resource in the cluster
632 | func (adm Admin) EnableResource(cluster string, resource string) error {
633 | // make sure the cluster is already setup
634 | if ok, err := adm.isClusterSetup(cluster); !ok || err != nil {
635 | return ErrClusterNotSetup
636 | }
637 |
638 | builder := KeyBuilder{cluster}
639 |
640 | isPath := builder.idealStates() + "/" + resource
641 |
642 | if exists, _, err := adm.zkClient.Exists(isPath); !exists || err != nil {
643 | if !exists {
644 | return ErrResourceNotExists
645 | }
646 | return err
647 | }
648 |
649 | // TODO: set the value at leaf node instead of the record level
650 | adm.zkClient.UpdateSimpleField(isPath, "HELIX_ENABLED", "true")
651 | return nil
652 | }
653 |
654 | // DisableResource disables the specified resource in the cluster.
655 | func (adm Admin) DisableResource(cluster string, resource string) error {
656 | // make sure the cluster is already setup
657 | if ok, err := adm.isClusterSetup(cluster); !ok || err != nil {
658 | return ErrClusterNotSetup
659 | }
660 |
661 | builder := KeyBuilder{cluster}
662 |
663 | isPath := builder.idealStates() + "/" + resource
664 |
665 | if exists, _, err := adm.zkClient.Exists(isPath); !exists || err != nil {
666 | if !exists {
667 | return ErrResourceNotExists
668 | }
669 |
670 | return err
671 | }
672 |
673 | // TODO: set the value at leaf node instead of the record level
674 | adm.zkClient.UpdateSimpleField(isPath, "HELIX_ENABLED", "false")
675 |
676 | return nil
677 | }
678 |
679 | // ListClusterInfo shows the existing resources and instances in the glaster
680 | func (adm Admin) ListClusterInfo(cluster string) (string, error) {
681 | // make sure the cluster is already setup
682 | if ok, err := adm.isClusterSetup(cluster); !ok || err != nil {
683 | return "", ErrClusterNotSetup
684 | }
685 |
686 | builder := KeyBuilder{cluster}
687 | isPath := builder.idealStates()
688 | instancesPath := builder.instances()
689 |
690 | resources, err := adm.zkClient.Children(isPath)
691 | if err != nil {
692 | return "", err
693 | }
694 |
695 | instances, err := adm.zkClient.Children(instancesPath)
696 | if err != nil {
697 | return "", err
698 | }
699 |
700 | var buffer bytes.Buffer
701 | buffer.WriteString("Existing resources in cluster " + cluster + ":\n")
702 |
703 | for _, r := range resources {
704 | buffer.WriteString(" " + r + "\n")
705 | }
706 |
707 | buffer.WriteString("\nInstances in cluster " + cluster + ":\n")
708 | for _, i := range instances {
709 | buffer.WriteString(" " + i + "\n")
710 | }
711 | return buffer.String(), nil
712 | }
713 |
714 | // ListClusters shows all Helix managed clusters in the connected zookeeper cluster
715 | func (adm Admin) ListClusters() (string, error) {
716 | var clusters []string
717 |
718 | children, err := adm.zkClient.Children("/")
719 | if err != nil {
720 | return "", err
721 | }
722 |
723 | for _, cluster := range children {
724 | if ok, err := adm.isClusterSetup(cluster); ok && err == nil {
725 | clusters = append(clusters, cluster)
726 | }
727 | }
728 |
729 | var buffer bytes.Buffer
730 | buffer.WriteString("Existing clusters: \n")
731 |
732 | for _, cluster := range clusters {
733 | buffer.WriteString(" " + cluster + "\n")
734 | }
735 | return buffer.String(), nil
736 | }
737 |
738 | // ListResources shows a list of resources managed by the helix cluster
739 | func (adm Admin) ListResources(cluster string) (string, error) {
740 | // make sure the cluster is already setup
741 | if ok, err := adm.isClusterSetup(cluster); !ok || err != nil {
742 | return "", ErrClusterNotSetup
743 | }
744 |
745 | builder := KeyBuilder{cluster}
746 | isPath := builder.idealStates()
747 | resources, err := adm.zkClient.Children(isPath)
748 | if err != nil {
749 | return "", err
750 | }
751 |
752 | var buffer bytes.Buffer
753 | buffer.WriteString("Existing resources in cluster " + cluster + ":\n")
754 |
755 | for _, r := range resources {
756 | buffer.WriteString(" " + r + "\n")
757 | }
758 |
759 | return buffer.String(), nil
760 | }
761 |
762 | // ListInstances shows a list of instances participating the cluster.
763 | func (adm Admin) ListInstances(cluster string) (string, error) {
764 | // make sure the cluster is already setup
765 | if ok, err := adm.isClusterSetup(cluster); !ok || err != nil {
766 | return "", ErrClusterNotSetup
767 | }
768 |
769 | builder := KeyBuilder{cluster}
770 | isPath := builder.instances()
771 | instances, err := adm.zkClient.Children(isPath)
772 | if err != nil {
773 | return "", err
774 | }
775 |
776 | var buffer bytes.Buffer
777 | buffer.WriteString(fmt.Sprintf("Existing instances in cluster %s:\n", cluster))
778 |
779 | for _, r := range instances {
780 | buffer.WriteString(" " + r + "\n")
781 | }
782 |
783 | return buffer.String(), nil
784 | }
785 |
786 | // ListInstanceInfo shows detailed information of an inspace in the helix cluster
787 | func (adm Admin) ListInstanceInfo(cluster string, instance string) (string, error) {
788 | // make sure the cluster is already setup
789 | if ok, err := adm.isClusterSetup(cluster); !ok || err != nil {
790 | return "", ErrClusterNotSetup
791 | }
792 |
793 | builder := &KeyBuilder{cluster}
794 | instanceCfg := builder.participantConfig(instance)
795 |
796 | if exists, _, err := adm.zkClient.Exists(instanceCfg); !exists || err != nil {
797 | if !exists {
798 | return "", ErrNodeNotExist
799 | }
800 | return "", err
801 | }
802 |
803 | accessor := newDataAccessor(adm.zkClient, builder)
804 | r, err := accessor.Msg(instanceCfg)
805 | if err != nil {
806 | return "", err
807 | }
808 | return r.String(), nil
809 | }
810 |
811 | // ListIdealState shows a list of ideal states for the cluster resource
812 | func (adm Admin) ListIdealState(cluster string, resource string) (*model.IdealState, error) {
813 | // make sure the cluster is already setup
814 | if ok, err := adm.isClusterSetup(cluster); !ok || err != nil {
815 | return nil, ErrClusterNotSetup
816 | }
817 |
818 | builder := &KeyBuilder{cluster}
819 | path := builder.idealStateForResource(resource)
820 |
821 | // check path exists
822 | if exists, _, err := adm.zkClient.Exists(path); !exists || err != nil {
823 | if !exists {
824 | return nil, ErrNodeNotExist
825 | }
826 | return nil, err
827 | }
828 |
829 | accessor := newDataAccessor(adm.zkClient, builder)
830 |
831 | return accessor.IdealState(resource)
832 | }
833 |
834 | // ListExternalView shows the externalviews for the cluster resource
835 | func (adm Admin) ListExternalView(cluster string, resource string) (*model.ExternalView, error) {
836 | // make sure the cluster is already setup
837 | if ok, err := adm.isClusterSetup(cluster); !ok || err != nil {
838 | return nil, ErrClusterNotSetup
839 | }
840 |
841 | builder := &KeyBuilder{cluster}
842 | path := builder.externalViewForResource(resource)
843 |
844 | // check path exists
845 | if exists, _, err := adm.zkClient.Exists(path); !exists || err != nil {
846 | if !exists {
847 | return nil, ErrNodeNotExist
848 | }
849 | return nil, err
850 | }
851 |
852 | accessor := newDataAccessor(adm.zkClient, builder)
853 |
854 | return accessor.ExternalView(resource)
855 | }
856 |
857 | // GetInstances prints out lists of instances
858 | func (adm Admin) GetInstances(cluster string) error {
859 | kb := KeyBuilder{cluster}
860 | instancesKey := kb.instances()
861 |
862 | data, _, err := adm.zkClient.Get(instancesKey)
863 | if err != nil {
864 | return err
865 | }
866 |
867 | for _, c := range data {
868 | fmt.Println(c)
869 | }
870 |
871 | return nil
872 | }
873 |
874 | // DropInstance removes a participating instance from the helix cluster
875 | func (adm Admin) DropInstance(cluster string, instance string) error {
876 | kb := KeyBuilder{cluster}
877 | instanceKey := kb.instance(instance)
878 | err := adm.zkClient.DeleteTree(instanceKey)
879 | if err != nil {
880 | return err
881 | }
882 |
883 | fmt.Printf("/%s/%s deleted from zookeeper.\n", cluster, instance)
884 |
885 | return err
886 | }
887 |
888 | func (adm Admin) isClusterSetup(cluster string) (bool, error) {
889 | keyBuilder := KeyBuilder{cluster}
890 |
891 | return adm.zkClient.ExistsAll(
892 | keyBuilder.cluster(),
893 | keyBuilder.idealStates(),
894 | keyBuilder.participantConfigs(),
895 | keyBuilder.propertyStore(),
896 | keyBuilder.liveInstances(),
897 | keyBuilder.instances(),
898 | keyBuilder.externalView(),
899 | keyBuilder.stateModelDefs(),
900 | )
901 | }
902 |
--------------------------------------------------------------------------------
/admin_test.go:
--------------------------------------------------------------------------------
1 | // Copyright (c) 2017 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 helix
22 |
23 | import (
24 | "encoding/json"
25 | "fmt"
26 | "strings"
27 | "testing"
28 | "time"
29 |
30 | "github.com/stretchr/testify/suite"
31 | "github.com/uber-go/go-helix/model"
32 | )
33 |
34 | type AdminTestSuite struct {
35 | BaseHelixTestSuite
36 | }
37 |
38 | func TestAdminTestSuite(t *testing.T) {
39 | suite.Run(t, &AdminTestSuite{})
40 | }
41 |
42 | func (s *AdminTestSuite) TestAddAndDropCluster() {
43 | t := s.T()
44 |
45 | // definitely a new cluster name by timestamp
46 | now := time.Now().Local()
47 | cluster := "AdminTest_TestAddAndDropCluster_" + now.Format("20060102150405")
48 |
49 | added := s.Admin.AddCluster(cluster, false)
50 | if !added {
51 | t.Error("Failed to add cluster: " + cluster)
52 | }
53 |
54 | // if cluster is already added, add it again and it should return false
55 | added = s.Admin.AddCluster(cluster, false)
56 | if added {
57 | t.Error("Should fail to add the same cluster")
58 | }
59 |
60 | //listClusters
61 | info, err := s.Admin.ListClusters()
62 | if err != nil || !strings.Contains(info, cluster) {
63 | t.Error("Expect OK")
64 | }
65 |
66 | s.Admin.DropCluster(cluster)
67 | info, err = s.Admin.ListClusters()
68 | if err != nil || strings.Contains(info, cluster) {
69 | t.Error("Expect dropped")
70 | }
71 | }
72 |
73 | func (s *AdminTestSuite) TestAddCluster() {
74 | now := time.Now().Local()
75 | cluster := "AdminTest_TestAddCluster" + now.Format("20060102150405")
76 |
77 | added := s.Admin.AddCluster(cluster, false)
78 | if added {
79 | defer s.Admin.DropCluster(cluster)
80 | }
81 |
82 | clusterInfo, err := s.Admin.ListClusterInfo(cluster)
83 | s.True(len(clusterInfo) > 0)
84 | s.NoError(err)
85 |
86 | // verify the data structure in zookeeper
87 | propertyStore := fmt.Sprintf("/%s/PROPERTYSTORE", cluster)
88 | s.verifyNodeExist(propertyStore)
89 |
90 | stateModelDefs := fmt.Sprintf("/%s/STATEMODELDEFS", cluster)
91 | s.verifyNodeExist(stateModelDefs)
92 | s.verifyChildrenCount(stateModelDefs, 6)
93 |
94 | instances := fmt.Sprintf("/%s/INSTANCES", cluster)
95 | s.verifyNodeExist(instances)
96 |
97 | instancesInfo, err := s.Admin.ListInstances(cluster)
98 | s.True(len(instancesInfo) > 0)
99 | s.NoError(err)
100 | err = s.Admin.GetInstances(cluster)
101 | s.NoError(err)
102 |
103 | configs := fmt.Sprintf("/%s/CONFIGS", cluster)
104 | s.verifyNodeExist(configs)
105 | s.verifyChildrenCount(configs, 3)
106 |
107 | idealStates := fmt.Sprintf("/%s/IDEALSTATES", cluster)
108 | s.verifyNodeExist(idealStates)
109 |
110 | externalView := fmt.Sprintf("/%s/EXTERNALVIEW", cluster)
111 | s.verifyNodeExist(externalView)
112 |
113 | liveInstances := fmt.Sprintf("/%s/LIVEINSTANCES", cluster)
114 | s.verifyNodeExist(liveInstances)
115 |
116 | controller := fmt.Sprintf("/%s/CONTROLLER", cluster)
117 | s.verifyNodeExist(controller)
118 | s.verifyChildrenCount(controller, 4)
119 | }
120 |
121 | func (s *AdminTestSuite) TestSetConfig() {
122 | t := s.T()
123 |
124 | now := time.Now().Local()
125 | cluster := "AdminTest_TestSetConfig_" + now.Format("20060102150405")
126 |
127 | added := s.Admin.AddCluster(cluster, false)
128 | if added {
129 | defer s.Admin.DropCluster(cluster)
130 | }
131 |
132 | property := map[string]string{
133 | _allowParticipantAutoJoinKey: "true",
134 | }
135 |
136 | err := s.Admin.SetConfig(cluster, "CLUSTER", property)
137 | s.NoError(err)
138 |
139 | prop, err := s.Admin.GetConfig(cluster, "CLUSTER", []string{_allowParticipantAutoJoinKey})
140 | s.NoError(err)
141 |
142 | if prop[_allowParticipantAutoJoinKey] != "true" {
143 | t.Error("allowParticipantAutoJoin config set/get failed")
144 | }
145 | err = s.Admin.SetConfig(cluster, "CONSTRAINT", property)
146 | s.NoError(err)
147 | prop, err = s.Admin.GetConfig(cluster, "CONSTRAINT", []string{})
148 | s.Equal(0, len(prop))
149 | s.NoError(err)
150 | err = s.Admin.SetConfig(cluster, "PARTICIPANT", property)
151 | s.NoError(err)
152 | prop, err = s.Admin.GetConfig(cluster, "PARTICIPANT", []string{})
153 | s.Equal(0, len(prop))
154 | s.NoError(err)
155 | err = s.Admin.SetConfig(cluster, "PARTITION", property)
156 | s.NoError(err)
157 | prop, err = s.Admin.GetConfig(cluster, "PARTITION", []string{})
158 | s.Equal(0, len(prop))
159 | s.NoError(err)
160 | err = s.Admin.SetConfig(cluster, "RESOURCE", property)
161 | s.NoError(err)
162 | prop, err = s.Admin.GetConfig(cluster, "RESOURCE", []string{})
163 | s.Equal(0, len(prop))
164 | s.NoError(err)
165 | }
166 |
167 | func (s *AdminTestSuite) TestAddDropNode() {
168 | t := s.T()
169 |
170 | // verify not able to add node before cluster is setup
171 | now := time.Now().Local()
172 | cluster := "AdminTest_TestAddDropNode_" + now.Format("20060102150405")
173 |
174 | node := "localhost_19932"
175 |
176 | // add node before adding cluster, expect fail
177 | if err := s.Admin.AddNode(cluster, node); err != ErrClusterNotSetup {
178 | t.Error("Must error out for AddNode if cluster not setup")
179 | }
180 |
181 | // now add the cluster and add the node again
182 | s.Admin.AddCluster(cluster, false)
183 | defer s.Admin.DropCluster(cluster)
184 |
185 | if err := s.Admin.AddNode(cluster, node); err != nil {
186 | t.Error("Should be able to add node")
187 | }
188 |
189 | // add the same node again, should expect error ErrNodeAlreadyExists
190 | if err := s.Admin.AddNode(cluster, node); err != ErrNodeAlreadyExists {
191 | t.Error("should not be able to add the same node")
192 | }
193 |
194 | // listInstanceInfo
195 | info, err := s.Admin.ListInstanceInfo(cluster, node)
196 | if err != nil || info == "" || !strings.Contains(info, node) {
197 | t.Error("expect OK")
198 | }
199 |
200 | // drop the node
201 | if err := s.Admin.DropNode(cluster, node); err != nil {
202 | t.Error("failed to drop cluster node")
203 | }
204 | // listInstanceInfo
205 | if _, err := s.Admin.ListInstanceInfo(cluster, node); err != ErrNodeNotExist {
206 | t.Error("expect OK")
207 | }
208 |
209 | // drop node again and we should see an error ErrNodeNotExist
210 | if err := s.Admin.DropNode(cluster, node); err != ErrNodeNotExist {
211 | t.Error("failed to see expected error ErrNodeNotExist")
212 | }
213 |
214 | // make sure the path does not exist in zookeeper
215 | s.verifyNodeNotExist(fmt.Sprintf("/%s/INSTANCES/%s", cluster, node))
216 | s.verifyNodeNotExist(fmt.Sprintf("/%s/CONFIGS/PARTICIPANT/%s", cluster, node))
217 | }
218 |
219 | func (s *AdminTestSuite) TestAddDropResource() {
220 | t := s.T()
221 |
222 | now := time.Now().Local()
223 | cluster := "AdminTest_TestAddResource_" + now.Format("20060102150405")
224 | resource := "resource"
225 |
226 | // expect error if cluster not setup
227 | err := s.Admin.AddResource(cluster, resource, 32, "MasterSlave")
228 | if err != ErrClusterNotSetup {
229 | t.Error("must setup cluster before addResource")
230 | }
231 | if err := s.Admin.DropResource(cluster, resource); err != ErrClusterNotSetup {
232 | t.Error("must setup cluster before addResource")
233 | }
234 | if _, err := s.Admin.ListResources(cluster); err != ErrClusterNotSetup {
235 | t.Error("must setup cluster")
236 | }
237 |
238 | s.Admin.AddCluster(cluster, false)
239 | defer s.Admin.DropCluster(cluster)
240 |
241 | // it is ok to dropResource before resource exists
242 | if err := s.Admin.DropResource(cluster, resource); err != nil {
243 | t.Error("expect OK")
244 | }
245 |
246 | // expect error if state model does not exist
247 | err = s.Admin.AddResource(cluster, resource, 32, "NotExistStateModel")
248 | if err != ErrStateModelDefNotExist {
249 | t.Error("must use valid state model")
250 | }
251 |
252 | // expect pass
253 | err = s.Admin.AddResource(cluster, resource, 32, "MasterSlave")
254 | if err != nil {
255 | t.Error("fail addResource")
256 | }
257 | // expect failure
258 | err = s.Admin.AddResource(cluster, resource, 32, "MasterSlave")
259 | s.Error(err)
260 |
261 | if info, err := s.Admin.ListResources(cluster); err != nil || info == "" {
262 | t.Error("expect OK")
263 | }
264 |
265 | kb := KeyBuilder{cluster}
266 | isPath := kb.idealStates() + "/resource"
267 | s.verifyNodeExist(isPath)
268 |
269 | if err := s.Admin.DropResource(cluster, resource); err != nil {
270 | t.Error("expect OK")
271 | }
272 |
273 | s.verifyNodeNotExist(isPath)
274 | }
275 |
276 | func (s *AdminTestSuite) TestEnableDisableResource() {
277 | t := s.T()
278 |
279 | now := time.Now().Local()
280 | cluster := "AdminTest_TestEnableDisableResource_" + now.Format("20060102150405")
281 | resource := "resource"
282 |
283 | // expect error if cluster not setup
284 | if err := s.Admin.EnableResource(cluster, resource); err != ErrClusterNotSetup {
285 | t.Error("must setup cluster before enableResource")
286 | }
287 | if err := s.Admin.DisableResource(cluster, resource); err != ErrClusterNotSetup {
288 | t.Error("must setup cluster before enableResource")
289 | }
290 |
291 | s.Admin.AddCluster(cluster, false)
292 | defer s.Admin.DropCluster(cluster)
293 |
294 | // expect error if resource not exist
295 | if err := s.Admin.EnableResource(cluster, resource); err != ErrResourceNotExists {
296 | t.Error("expect ErrResourceNotExists")
297 | }
298 | if err := s.Admin.DisableResource(cluster, resource); err != ErrResourceNotExists {
299 | t.Error("expect ErrResourceNotExists")
300 | }
301 | if err := s.Admin.AddResource(cluster, "resource", 32, "MasterSlave"); err != nil {
302 | t.Error("fail addResource")
303 | }
304 | if err := s.Admin.EnableResource(cluster, resource); err != nil {
305 | // expect error if resource not exist
306 | t.Error("expect OK")
307 | }
308 | }
309 |
310 | func (s *AdminTestSuite) TestListExternalView() {
311 | t := s.T()
312 |
313 | now := time.Now().Local()
314 | cluster := "AdminTest_TestListExternalView_" + now.Format("20060102150405")
315 | resource := "resource"
316 |
317 | // expect error if cluster not setup
318 | if _, err := s.Admin.ListExternalView(cluster, resource); err != ErrClusterNotSetup {
319 | t.Error("must setup cluster before ListExternalView")
320 | }
321 |
322 | s.Admin.AddCluster(cluster, false)
323 | defer s.Admin.DropCluster(cluster)
324 |
325 | // fail when resource doesn't exist
326 | if _, err := s.Admin.ListExternalView(cluster, resource); err != ErrNodeNotExist {
327 | t.Error("must setup resource before ListExternalView")
328 | }
329 | // expect pass
330 | if err := s.Admin.AddResource(cluster, "resource", 32, "MasterSlave"); err != nil {
331 | t.Error("fail addResource")
332 | }
333 | // should fail, doesn't create externalView by default
334 | if _, err := s.Admin.ListExternalView(cluster, resource); err != ErrNodeNotExist {
335 | t.Error("must setup resource externalView before ListExternalView")
336 | }
337 |
338 | // create externalview
339 | externalView := fmt.Sprintf("/%s/EXTERNALVIEW/%s", cluster, resource)
340 | m := model.NewRecord("resource")
341 | m.SetIntField("NUM_PARTITIONS", 32)
342 | data, err := json.Marshal(m)
343 | if err != nil {
344 | t.Error("expect OK")
345 | }
346 | s.Admin.zkClient.CreateDataWithPath(externalView, data)
347 |
348 | res, err := s.Admin.ListExternalView(cluster, resource)
349 | if err != nil {
350 | t.Error("expect OK")
351 | }
352 | if res.ZNRecord.ID != "resource" || res.GetNumPartitions() != 32 {
353 | t.Error("expect read model OK")
354 | }
355 | }
356 |
357 | func (s *AdminTestSuite) TestListIdealState() {
358 | t := s.T()
359 |
360 | now := time.Now().Local()
361 | cluster := "AdminTest_TestListIdealState_" + now.Format("20060102150405")
362 | resource := "resource"
363 |
364 | // expect error if cluster not setup
365 | if _, err := s.Admin.ListIdealState(cluster, resource); err != ErrClusterNotSetup {
366 | t.Error("must setup cluster before ListIdealState")
367 | }
368 |
369 | s.Admin.AddCluster(cluster, false)
370 | defer s.Admin.DropCluster(cluster)
371 |
372 | // fail when resource doesn't exist
373 | if _, err := s.Admin.ListIdealState(cluster, resource); err != ErrNodeNotExist {
374 | t.Error("must setup resource before ListIdealState")
375 | }
376 | // expect pass
377 | if err := s.Admin.AddResource(cluster, "resource", 32, "MasterSlave"); err != nil {
378 | t.Error("fail addResource")
379 | }
380 |
381 | res, err := s.Admin.ListIdealState(cluster, resource)
382 | if err != nil {
383 | t.Error("expect OK")
384 | }
385 | if res.ZNRecord.ID != "resource" || res.GetNumPartitions() != 32 {
386 | t.Error("expect read model OK")
387 | }
388 | }
389 |
390 | func (s *AdminTestSuite) verifyNodeExist(path string) {
391 | if exists, _, err := s.Admin.zkClient.Exists(path); err != nil || !exists {
392 | s.T().Error("failed verifyNodeExist")
393 | }
394 | }
395 |
396 | func (s *AdminTestSuite) verifyNodeNotExist(path string) {
397 | if exists, _, err := s.Admin.zkClient.Exists(path); err != nil || exists {
398 | s.T().Error("failed verifyNotNotExist")
399 | s.T().FailNow()
400 | }
401 | }
402 |
403 | func (s *AdminTestSuite) verifyChildrenCount(path string, count int) {
404 | children, err := s.Admin.zkClient.Children(path)
405 | if err != nil {
406 | s.T().FailNow()
407 | }
408 |
409 | if len(children) != count {
410 | s.T().Errorf("Node %s should have %d children, "+
411 | "but only have %d children", path, count, len(children))
412 | }
413 | }
414 |
--------------------------------------------------------------------------------
/check_license.sh:
--------------------------------------------------------------------------------
1 | #!/bin/bash -e
2 |
3 | ERROR_COUNT=0
4 | while read -r file
5 | do
6 | case "$(head -1 "${file}")" in
7 | *"Copyright (c) "*" Uber Technologies, Inc.")
8 | # everything's cool
9 | ;;
10 | *)
11 | echo "$file is missing license header."
12 | (( ERROR_COUNT++ ))
13 | ;;
14 | esac
15 | done < <(git ls-files "*\.go")
16 |
17 | exit $ERROR_COUNT
18 |
--------------------------------------------------------------------------------
/constants.go:
--------------------------------------------------------------------------------
1 | // Copyright (c) 2017 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 helix
22 |
23 | // Helix constants
24 | const (
25 | StateModelNameOnlineOffline = "OnlineOffline"
26 |
27 | StateModelStateOnline = "ONLINE"
28 | StateModelStateOffline = "OFFLINE"
29 | StateModelStateDropped = "DROPPED"
30 |
31 | TargetController = "CONTROLLER"
32 |
33 | MsgTypeStateTransition = "STATE_TRANSITION"
34 | MsgTypeNoop = "NO_OP"
35 | )
36 |
--------------------------------------------------------------------------------
/cover.sh:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env bash
2 |
3 | set -e
4 | echo "" > coverage.txt
5 |
6 | for d in $(go list $@); do
7 | go test -race -coverprofile=profile.log $d
8 | if [ -f profile.log ]; then
9 | cat profile.log >> coverage.txt
10 | rm profile.log
11 | fi
12 | done
13 |
--------------------------------------------------------------------------------
/data_accessor.go:
--------------------------------------------------------------------------------
1 | // Copyright (c) 2017 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 helix
22 |
23 | import (
24 | "github.com/pkg/errors"
25 | "github.com/samuel/go-zookeeper/zk"
26 | "github.com/uber-go/go-helix/model"
27 | uzk "github.com/uber-go/go-helix/zk"
28 | )
29 |
30 | var (
31 | // ErrorNilUpdatedData returns when updateFn returns nil *model.ZNRecord without error
32 | ErrorNilUpdatedData = errors.New("data accessor: updated data is nil")
33 | )
34 |
35 | // updateFn gets the old data (nil if data does not exist) and return the updated data
36 | // return data is expected to be non-nil
37 | type updateFn func(data *model.ZNRecord) (*model.ZNRecord, error)
38 |
39 | // DataAccessor helps to interact with Helix Data Types like IdealState, LiveInstance, Message
40 | // it mirrors org.apache.helix.HelixDataAccessor
41 | type DataAccessor struct {
42 | zkClient *uzk.Client
43 | keyBuilder *KeyBuilder
44 | }
45 |
46 | // newDataAccessor creates new DataAccessor with Zookeeper client
47 | func newDataAccessor(zkClient *uzk.Client, keyBuilder *KeyBuilder) *DataAccessor {
48 | return &DataAccessor{zkClient: zkClient, keyBuilder: keyBuilder}
49 | }
50 |
51 | // Msg helps get Helix property with type Message
52 | func (a *DataAccessor) Msg(path string) (*model.Message, error) {
53 | record, err := a.zkClient.GetRecordFromPath(path)
54 | if err != nil {
55 | return nil, err
56 | }
57 | return &model.Message{ZNRecord: *record}, nil
58 | }
59 |
60 | // InstanceConfig helps get Helix property with type Message
61 | func (a *DataAccessor) InstanceConfig(path string) (*model.InstanceConfig, error) {
62 | record, err := a.zkClient.GetRecordFromPath(path)
63 | if err != nil {
64 | return nil, err
65 | }
66 | return &model.InstanceConfig{ZNRecord: *record}, nil
67 | }
68 |
69 | // IdealState helps get Helix property with type IdealState
70 | func (a *DataAccessor) IdealState(resourceName string) (*model.IdealState, error) {
71 | path := a.keyBuilder.idealStateForResource(resourceName)
72 | record, err := a.zkClient.GetRecordFromPath(path)
73 | if err != nil {
74 | return nil, err
75 | }
76 | return &model.IdealState{ZNRecord: *record}, nil
77 | }
78 |
79 | // ExternalView helps get Helix property with type ExternalView
80 | func (a *DataAccessor) ExternalView(resourceName string) (*model.ExternalView, error) {
81 | path := a.keyBuilder.externalViewForResource(resourceName)
82 | record, err := a.zkClient.GetRecordFromPath(path)
83 | if err != nil {
84 | return nil, err
85 | }
86 | return &model.ExternalView{ZNRecord: *record}, nil
87 | }
88 |
89 | // CurrentState helps get Helix property with type CurrentState
90 | func (a *DataAccessor) CurrentState(instanceName, session, resourceName string,
91 | ) (*model.CurrentState, error) {
92 | path := a.keyBuilder.currentStateForResource(instanceName, session, resourceName)
93 | record, err := a.zkClient.GetRecordFromPath(path)
94 | if err != nil {
95 | return nil, err
96 | }
97 | return &model.CurrentState{ZNRecord: *record}, nil
98 | }
99 |
100 | // LiveInstance returns Helix property with type LiveInstance
101 | func (a *DataAccessor) LiveInstance(instanceName string) (*model.LiveInstance, error) {
102 | path := a.keyBuilder.liveInstance(instanceName)
103 | record, err := a.zkClient.GetRecordFromPath(path)
104 | if err != nil {
105 | return nil, err
106 | }
107 | return &model.LiveInstance{ZNRecord: *record}, nil
108 | }
109 |
110 | // StateModelDef helps get Helix property with type StateModelDef
111 | func (a *DataAccessor) StateModelDef(stateModel string) (*model.StateModelDef, error) {
112 | path := a.keyBuilder.stateModelDef(stateModel)
113 | record, err := a.zkClient.GetRecordFromPath(path)
114 | if err != nil {
115 | return nil, err
116 | }
117 | return &model.StateModelDef{ZNRecord: *record}, nil
118 | }
119 |
120 | // updateData would update the data in path with updateFn
121 | // if path does not exist, updateData would create it
122 | // and updateFn would have a nil *model.ZNRecord as input
123 | func (a *DataAccessor) updateData(path string, update updateFn) error {
124 | var err error
125 | var record *model.ZNRecord
126 | for {
127 | nodeExists := true
128 | record, err = a.zkClient.GetRecordFromPath(path)
129 | cause := errors.Cause(err)
130 | if cause == zk.ErrNoNode {
131 | nodeExists = false
132 | } else if cause != nil {
133 | return err
134 | }
135 |
136 | record, err = update(record)
137 | if err != nil {
138 | return err
139 | }
140 | if record == nil {
141 | return ErrorNilUpdatedData
142 | }
143 | if nodeExists {
144 | err = a.setData(path, *record, record.Version)
145 | } else {
146 | err = a.createData(path, *record)
147 | }
148 |
149 | if cause := errors.Cause(err); cause != nil &&
150 | (cause != zk.ErrBadVersion || cause != zk.ErrNoNode) {
151 | return err
152 | }
153 |
154 | // retry for ErrBadVersion or ErrNoNode
155 | if err == nil {
156 | return nil
157 | }
158 | }
159 | }
160 |
161 | func (a *DataAccessor) createData(path string, data model.ZNRecord) error {
162 | serialized, err := data.Marshal()
163 | if err != nil {
164 | return err
165 | }
166 | return a.zkClient.CreateDataWithPath(path, serialized)
167 | }
168 |
169 | func (a *DataAccessor) setData(path string, data model.ZNRecord, version int32) error {
170 | serialized, err := data.Marshal()
171 | if err != nil {
172 | return err
173 | }
174 | return a.zkClient.SetDataForPath(path, serialized, version)
175 | }
176 |
177 | // createMsg creates a new message
178 | func (a *DataAccessor) createMsg(path string, msg *model.Message) error {
179 | return a.createData(path, msg.ZNRecord)
180 | }
181 |
182 | // CreateParticipantMsg creates a message for the participant and helps
183 | // testing the participant
184 | func (a *DataAccessor) CreateParticipantMsg(instanceName string, msg *model.Message) error {
185 | path := a.keyBuilder.participantMsg(instanceName, msg.ID)
186 | return a.createData(path, msg.ZNRecord)
187 | }
188 |
189 | func (a *DataAccessor) setMsg(path string, msg *model.Message) error {
190 | return a.setData(path, msg.ZNRecord, msg.Version)
191 | }
192 |
193 | func (a *DataAccessor) createCurrentState(path string, state *model.CurrentState) error {
194 | return a.createData(path, state.ZNRecord)
195 | }
196 |
197 | func (a *DataAccessor) createInstanceConfig(path string, config *model.InstanceConfig) error {
198 | return a.createData(path, config.ZNRecord)
199 | }
200 |
--------------------------------------------------------------------------------
/data_accessor_test.go:
--------------------------------------------------------------------------------
1 | // Copyright (c) 2017 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 helix
22 |
23 | import (
24 | "fmt"
25 | "math/rand"
26 | "testing"
27 |
28 | "github.com/stretchr/testify/suite"
29 | "github.com/uber-go/go-helix/model"
30 | )
31 |
32 | var (
33 | _accessorTestKeyBuilder = &KeyBuilder{TestClusterName}
34 | )
35 |
36 | type DataAccessorTestSuite struct {
37 | BaseHelixTestSuite
38 | }
39 |
40 | func TestDataAccessorTestSuite(t *testing.T) {
41 | suite.Run(t, &DataAccessorTestSuite{})
42 | }
43 |
44 | func (s *DataAccessorTestSuite) TestMsg() {
45 | client := s.CreateAndConnectClient()
46 | defer client.Disconnect()
47 | accessor := newDataAccessor(client, _accessorTestKeyBuilder)
48 | path := fmt.Sprintf("/test_path/msg/%s", CreateRandomString())
49 | err := accessor.createMsg(path, model.NewMsg("test_id"))
50 | s.NoError(err)
51 | msg, err := accessor.Msg(path)
52 | s.NoError(err)
53 | msg.SetMsgState(model.MessageStateRead)
54 | err = accessor.setMsg(path, msg)
55 | s.NoError(err)
56 | s.Equal(model.MessageStateRead, msg.GetMsgState())
57 | }
58 |
59 | func (s *DataAccessorTestSuite) TestCurrentState() {
60 | client := s.CreateAndConnectClient()
61 | defer client.Disconnect()
62 | accessor := newDataAccessor(client, _accessorTestKeyBuilder)
63 | session := CreateRandomString()
64 | resource := CreateRandomString()
65 | path := _accessorTestKeyBuilder.currentStateForResource(testInstanceName, session, resource)
66 | msg := model.NewMsg(CreateRandomString())
67 | msg.SetStateModelDef(StateModelNameOnlineOffline)
68 | state := model.NewCurrentStateFromMsg(msg, resource, session)
69 | state.SetMapField("partition_1", model.FieldKeyCurrentState, StateModelStateOnline)
70 | state.SetMapField("partition_2", model.FieldKeyCurrentState, StateModelStateOffline)
71 | err := accessor.createCurrentState(path, state)
72 | s.NoError(err)
73 | state, err = accessor.CurrentState(testInstanceName, session, resource)
74 | s.NoError(err)
75 | s.Equal(session, state.GetSessionID())
76 | partitionStateMap := state.GetPartitionStateMap()
77 | s.Len(partitionStateMap, 2)
78 | s.assertCurrentState(state, "partition_1", StateModelStateOnline)
79 | s.assertCurrentState(state, "partition_2", StateModelStateOffline)
80 | s.Equal(state.GetStateModelDef(), StateModelNameOnlineOffline)
81 | }
82 |
83 | func (s *DataAccessorTestSuite) TestUpdateExistingCurrentState() {
84 | client := s.CreateAndConnectClient()
85 | defer client.Disconnect()
86 | accessor := newDataAccessor(client, _accessorTestKeyBuilder)
87 | session := CreateRandomString()
88 | resource := CreateRandomString()
89 | path := _accessorTestKeyBuilder.currentStateForResource(testInstanceName, session, resource)
90 | msg := model.NewMsg(CreateRandomString())
91 | msg.SetStateModelDef(StateModelNameOnlineOffline)
92 | state := model.NewCurrentStateFromMsg(msg, resource, session)
93 | state.SetMapField("partition_1", model.FieldKeyCurrentState, StateModelStateOnline)
94 | state.SetMapField("partition_2", model.FieldKeyCurrentState, StateModelStateOffline)
95 | err := accessor.createCurrentState(path, state)
96 | s.NoError(err)
97 |
98 | err = accessor.updateData(path, func(data *model.ZNRecord) (*model.ZNRecord, error) {
99 | currentState := &model.CurrentState{ZNRecord: *data}
100 | currentState.SetState("partition_1", StateModelStateOffline)
101 | currentState.SetState("partition_2", StateModelStateOnline)
102 | return ¤tState.ZNRecord, nil
103 | })
104 | s.NoError(err)
105 |
106 | state, err = accessor.CurrentState(testInstanceName, session, resource)
107 | s.NoError(err)
108 | s.assertCurrentState(state, "partition_1", StateModelStateOffline)
109 | s.assertCurrentState(state, "partition_2", StateModelStateOnline)
110 | s.Equal(state.GetStateModelDef(), StateModelNameOnlineOffline)
111 | }
112 |
113 | func (s *DataAccessorTestSuite) TestUpdateNonExistingCurrentState() {
114 | client := s.CreateAndConnectClient()
115 | defer client.Disconnect()
116 | accessor := newDataAccessor(client, _accessorTestKeyBuilder)
117 | session := CreateRandomString()
118 | resource := CreateRandomString()
119 | path := _accessorTestKeyBuilder.currentStateForResource(testInstanceName, session, resource)
120 |
121 | err := accessor.updateData(path, func(data *model.ZNRecord) (*model.ZNRecord, error) {
122 | if data != nil {
123 | s.Fail("data is expected to be nil")
124 | }
125 | currentState := &model.CurrentState{ZNRecord: *model.NewRecord(resource)}
126 | currentState.SetState("partition_1", StateModelStateOffline)
127 | currentState.SetState("partition_2", StateModelStateOnline)
128 | return ¤tState.ZNRecord, nil
129 | })
130 | s.NoError(err)
131 |
132 | state, err := accessor.CurrentState(testInstanceName, session, resource)
133 | s.NoError(err)
134 | s.assertCurrentState(state, "partition_1", StateModelStateOffline)
135 | s.assertCurrentState(state, "partition_2", StateModelStateOnline)
136 | }
137 |
138 | func (s *DataAccessorTestSuite) TestInstanceConfig() {
139 | client := s.CreateAndConnectClient()
140 | defer client.Disconnect()
141 | accessor := newDataAccessor(client, _accessorTestKeyBuilder)
142 | path := fmt.Sprintf("/test_path/instance_config/%s", CreateRandomString())
143 | config := model.NewInstanceConfig(testInstanceName)
144 | err := accessor.createInstanceConfig(path, config)
145 | s.NoError(err)
146 | config, err = accessor.InstanceConfig(path)
147 | s.False(config.GetEnabled())
148 | }
149 |
150 | func (s *DataAccessorTestSuite) TestGetStateModelDef() {
151 | keyBuilder := &KeyBuilder{TestClusterName}
152 | client := s.CreateAndConnectClient()
153 | defer client.Disconnect()
154 | accessor := newDataAccessor(client, keyBuilder)
155 |
156 | onlineOfflineModel, err := accessor.StateModelDef(StateModelNameOnlineOffline)
157 | s.NoError(err)
158 | s.Equal(onlineOfflineModel.GetInitialState(), StateModelStateOffline)
159 |
160 | masterSlaveMode, err := accessor.StateModelDef("MasterSlave")
161 | s.NoError(err)
162 | s.Equal(masterSlaveMode.GetInitialState(), StateModelStateOffline)
163 |
164 | leaderStandbyModel, err := accessor.StateModelDef("LeaderStandby")
165 | s.NoError(err)
166 | s.Equal(leaderStandbyModel.GetInitialState(), StateModelStateOffline)
167 | }
168 |
169 | func (s *DataAccessorTestSuite) TestIdealState() {
170 | cluster, resource := CreateRandomString(), CreateRandomString()
171 | admin, err := NewAdmin(s.ZkConnectString)
172 | s.NoError(err)
173 | addedCluster := admin.AddCluster(cluster, false)
174 | s.True(addedCluster)
175 | numPartitions := rand.Int()
176 | err = admin.AddResource(cluster, resource, numPartitions, StateModelNameOnlineOffline)
177 | s.NoError(err)
178 |
179 | keyBuilder := &KeyBuilder{cluster}
180 | client := s.CreateAndConnectClient()
181 | defer client.Disconnect()
182 | accessor := newDataAccessor(client, keyBuilder)
183 | idealState, err := accessor.IdealState(resource)
184 | s.NoError(err)
185 | s.Equal(numPartitions, idealState.GetNumPartitions())
186 | }
187 |
188 | func (s *DataAccessorTestSuite) TestLiveInstance() {
189 | p, _ := s.createParticipantAndConnect()
190 | s.NotNil(p)
191 |
192 | keyBuilder := &KeyBuilder{TestClusterName}
193 | client := s.CreateAndConnectClient()
194 | defer client.Disconnect()
195 | accessor := newDataAccessor(client, keyBuilder)
196 |
197 | liveInstance, err := accessor.LiveInstance(p.InstanceName())
198 | s.NoError(err)
199 | s.Equal(liveInstance.GetSessionID(), p.zkClient.GetSessionID())
200 | }
201 |
202 | func (s *DataAccessorTestSuite) assertCurrentState(
203 | state *model.CurrentState,
204 | partition string,
205 | expectedState string,
206 | ) {
207 | s.Equal(state.GetPartitionStateMap()[partition], expectedState)
208 | s.Equal(state.GetState(partition), expectedState)
209 | }
210 |
--------------------------------------------------------------------------------
/glide.lock:
--------------------------------------------------------------------------------
1 | hash: 66fcadd1753f7ee63c624c0550b46a2867defe89563adc6b2811b947792bce64
2 | updated: 2017-12-15T15:10:07.964694982-08:00
3 | imports:
4 | - name: github.com/davecgh/go-spew
5 | version: 04cdfd42973bb9c8589fd6a731800cf222fde1a9
6 | subpackages:
7 | - spew
8 | - name: github.com/facebookgo/clock
9 | version: 600d898af40aa09a7a93ecb9265d87b0504b6f03
10 | - name: github.com/pkg/errors
11 | version: 8842a6e0cc595d1cc9d931f6c875883967280e32
12 | - name: github.com/pmezard/go-difflib
13 | version: d8ed2627bdf02c080bf22230dbb337003b7aba2d
14 | subpackages:
15 | - difflib
16 | - name: github.com/samuel/go-zookeeper
17 | version: 471cd4e61d7a78ece1791fa5faa0345dc8c7d5a5
18 | subpackages:
19 | - zk
20 | - name: github.com/stretchr/testify
21 | version: 2aa2c176b9dab406a6970f6a55f513e8a8c8b18f
22 | subpackages:
23 | - assert
24 | - require
25 | - suite
26 | - name: github.com/uber-go/tally
27 | version: 6c4631652c6aab57c64f65c2e0aaec2e9aae3a64
28 | - name: go.uber.org/atomic
29 | version: 8474b86a5a6f79c443ce4b2992817ff32cf208b8
30 | - name: go.uber.org/multierr
31 | version: 3c4937480c32f4c13a875a1829af76c98ca3d40a
32 | - name: go.uber.org/zap
33 | version: f85c78b1dd998214c5f2138155b320a4a43fbe36
34 | subpackages:
35 | - buffer
36 | - internal/bufferpool
37 | - internal/color
38 | - internal/exit
39 | - zapcore
40 | testImports:
41 | - name: github.com/golang/lint
42 | version: f635bddafc7154957bd70209ee858a4b97e64a9b
43 | subpackages:
44 | - golint
45 | - name: golang.org/x/tools
46 | version: 64890f4e2b733655fee5077a5435a8812404c3a3
47 | subpackages:
48 | - go/gcexportdata
49 |
--------------------------------------------------------------------------------
/glide.yaml:
--------------------------------------------------------------------------------
1 | package: github.com/uber-go/go-helix
2 | license: MIT
3 | import:
4 | - package: github.com/samuel/go-zookeeper
5 | - package: github.com/pkg/errors
6 | - package: github.com/uber-go/tally
7 | - package: go.uber.org/zap
8 | testImport:
9 | - package: github.com/golang/lint
10 | subpackages:
11 | - golint
12 | - package: golang.org/x/tools
13 | subpackages:
14 | - go/gcexportdata
15 | - package: github.com/stretchr/testify
16 | subpackages:
17 | - mock
18 |
--------------------------------------------------------------------------------
/key_builder.go:
--------------------------------------------------------------------------------
1 | // Copyright (c) 2017 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 helix
22 |
23 | import (
24 | "fmt"
25 | )
26 |
27 | // KeyBuilder generates Zookeeper paths
28 | // Mirrors org.apache.helix.PropertyKey#Builder
29 | type KeyBuilder struct {
30 | clusterName string
31 | }
32 |
33 | func (b *KeyBuilder) cluster() string {
34 | return fmt.Sprintf("/%s", b.clusterName)
35 | }
36 |
37 | func (b *KeyBuilder) clusterConfig() string {
38 | return fmt.Sprintf("/%s/CONFIGS/CLUSTER/%s", b.clusterName, b.clusterName)
39 | }
40 |
41 | func (b *KeyBuilder) controller() string {
42 | return fmt.Sprintf("/%s/CONTROLLER", b.clusterName)
43 | }
44 |
45 | func (b *KeyBuilder) controllerMessages() string {
46 | return fmt.Sprintf("/%s/CONTROLLER/MESSAGES", b.clusterName)
47 | }
48 |
49 | func (b *KeyBuilder) controllerErrors() string {
50 | return fmt.Sprintf("/%s/CONTROLLER/ERRORS", b.clusterName)
51 | }
52 |
53 | func (b *KeyBuilder) controllerStatusUpdates() string {
54 | return fmt.Sprintf("/%s/CONTROLLER/STATUSUPDATES", b.clusterName)
55 | }
56 |
57 | func (b *KeyBuilder) controllerHistory() string {
58 | return fmt.Sprintf("/%s/CONTROLLER/HISTORY", b.clusterName)
59 | }
60 |
61 | func (b *KeyBuilder) externalView() string {
62 | return fmt.Sprintf("/%s/EXTERNALVIEW", b.clusterName)
63 | }
64 |
65 | func (b *KeyBuilder) externalViewForResource(resource string) string {
66 | return fmt.Sprintf("/%s/EXTERNALVIEW/%s", b.clusterName, resource)
67 | }
68 |
69 | func (b *KeyBuilder) propertyStore() string {
70 | return fmt.Sprintf("/%s/PROPERTYSTORE", b.clusterName)
71 | }
72 |
73 | func (b *KeyBuilder) idealStates() string {
74 | return fmt.Sprintf("/%s/IDEALSTATES", b.clusterName)
75 | }
76 |
77 | // IdealStateForResource returns path for ideal state of a resource
78 | func (b *KeyBuilder) idealStateForResource(resource string) string {
79 | return fmt.Sprintf("/%s/IDEALSTATES/%s", b.clusterName, resource)
80 | }
81 |
82 | func (b *KeyBuilder) resourceConfigs() string {
83 | return fmt.Sprintf("/%s/CONFIGS/RESOURCE", b.clusterName)
84 | }
85 |
86 | func (b *KeyBuilder) resourceConfig(resource string) string {
87 | return fmt.Sprintf("/%s/CONFIGS/RESOURCE/%s", b.clusterName, resource)
88 | }
89 |
90 | func (b *KeyBuilder) participantConfigs() string {
91 | return fmt.Sprintf("/%s/CONFIGS/PARTICIPANT", b.clusterName)
92 | }
93 |
94 | func (b *KeyBuilder) participantConfig(participantID string) string {
95 | return fmt.Sprintf("/%s/CONFIGS/PARTICIPANT/%s", b.clusterName, participantID)
96 | }
97 |
98 | func (b *KeyBuilder) liveInstances() string {
99 | return fmt.Sprintf("/%s/LIVEINSTANCES", b.clusterName)
100 | }
101 |
102 | func (b *KeyBuilder) instances() string {
103 | return fmt.Sprintf("/%s/INSTANCES", b.clusterName)
104 | }
105 |
106 | func (b *KeyBuilder) instance(participantID string) string {
107 | return fmt.Sprintf("/%s/INSTANCES/%s", b.clusterName, participantID)
108 | }
109 |
110 | func (b *KeyBuilder) liveInstance(partipantID string) string {
111 | return fmt.Sprintf("/%s/LIVEINSTANCES/%s", b.clusterName, partipantID)
112 | }
113 |
114 | func (b *KeyBuilder) currentStates(participantID string) string {
115 | return fmt.Sprintf("/%s/INSTANCES/%s/CURRENTSTATES", b.clusterName, participantID)
116 | }
117 |
118 | func (b *KeyBuilder) currentStatesForSession(participantID string, sessionID string) string {
119 | return fmt.Sprintf("/%s/INSTANCES/%s/CURRENTSTATES/%s", b.clusterName, participantID, sessionID)
120 | }
121 |
122 | func (b *KeyBuilder) currentStateForResource(
123 | participantID string, sessionID string, resourceID string) string {
124 | return fmt.Sprintf(
125 | "/%s/INSTANCES/%s/CURRENTSTATES/%s/%s", b.clusterName, participantID, sessionID, resourceID)
126 | }
127 |
128 | func (b *KeyBuilder) errorsR(participantID string) string {
129 | return fmt.Sprintf("/%s/INSTANCES/%s/ERRORS", b.clusterName, participantID)
130 | }
131 |
132 | func (b *KeyBuilder) errors(participantID string, sessionID string, resourceID string) string {
133 | return fmt.Sprintf(
134 | "/%s/INSTANCES/%s/ERRORS/%s/%s", b.clusterName, participantID, sessionID, resourceID)
135 | }
136 |
137 | func (b *KeyBuilder) healthReport(participantID string) string {
138 | return fmt.Sprintf("/%s/INSTANCES/%s/HEALTHREPORT", b.clusterName, participantID)
139 | }
140 |
141 | func (b *KeyBuilder) statusUpdates(participantID string) string {
142 | return fmt.Sprintf("/%s/INSTANCES/%s/STATUSUPDATES", b.clusterName, participantID)
143 | }
144 |
145 | func (b *KeyBuilder) stateModelDefs() string {
146 | return fmt.Sprintf("/%s/STATEMODELDEFS", b.clusterName)
147 | }
148 |
149 | func (b *KeyBuilder) stateModelDef(stateModel string) string {
150 | return fmt.Sprintf("/%s/STATEMODELDEFS/%s", b.clusterName, stateModel)
151 | }
152 |
153 | func (b *KeyBuilder) participantMessages(participantID string) string {
154 | return fmt.Sprintf("/%s/INSTANCES/%s/MESSAGES", b.clusterName, participantID)
155 | }
156 |
157 | func (b *KeyBuilder) participantMsg(participantID string, messageID string) string {
158 | return fmt.Sprintf("/%s/INSTANCES/%s/MESSAGES/%s", b.clusterName, participantID, messageID)
159 | }
160 |
161 | // IdealStateKey returns path for ideal state of a given cluster and resource
162 | func IdealStateKey(clusterName string, resourceName string) string {
163 | return fmt.Sprintf("/%s/IDEALSTATES/%s", clusterName, resourceName)
164 | }
165 |
--------------------------------------------------------------------------------
/model/constants.go:
--------------------------------------------------------------------------------
1 | // Copyright (c) 2017 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 model
22 |
23 | const (
24 | _helixVersion = "helix-0.6.8"
25 | _defaultStateModelFactoryName = "DEFAULT"
26 | )
27 |
28 | // Field keys commonly shared among Helix models
29 | const (
30 | FieldKeyBucketSize = "BUCKET_SIZE"
31 | FieldKeySessionID = "SESSION_ID"
32 | FieldKeyBatchMsgMode = "BATCH_MESSAGE_MODE"
33 | FieldKeyStateModelFactoryName = "STATE_MODEL_FACTORY_NAME"
34 | )
35 |
36 | // Field keys used by Helix message
37 | const (
38 | FieldKeyStateModelDef = "STATE_MODEL_DEF"
39 | FieldKeyTargetSessionID = "TGT_SESSION_ID"
40 | FieldKeyTargetName = "TGT_NAME"
41 | FieldKeyPartitionName = "PARTITION_NAME"
42 | FieldKeyFromState = "FROM_STATE"
43 | FieldKeyToState = "TO_STATE"
44 | FieldKeyCurrentState = "CURRENT_STATE"
45 | FieldKeyParentMsgID = "PARENT_MSG_ID"
46 | FieldKeyMsgState = "MSG_STATE"
47 | FieldKeyMsgType = "MSG_TYPE"
48 | FieldKeyResourceName = "RESOURCE_NAME"
49 | FieldKeyCreateTimestamp = "CREATE_TIMESTAMP"
50 | FieldKeyExecuteStartTimestamp = "EXECUTE_START_TIMESTAMP"
51 | )
52 |
53 | // Field keys used by the ideal state
54 | const (
55 | FieldKeyNumPartitions = "NUM_PARTITIONS"
56 | )
57 |
58 | // Field keys used by instance config
59 | const (
60 | FieldKeyHelixHost = "HELIX_HOST"
61 | FieldKeyHelixPort = "HELIX_PORT"
62 | FieldKeyHelixEnabled = "HELIX_ENABLED"
63 | )
64 |
65 | // Field keys used by live instance
66 | const (
67 | FieldKeyHelixVersion = "HELIX_VERSION"
68 | FieldKeyLiveInstance = "LIVE_INSTANCE"
69 | )
70 |
71 | // Field keys used by state model def
72 | const (
73 | FieldKeyInitialState = "INITIAL_STATE"
74 | )
75 |
--------------------------------------------------------------------------------
/model/current_state.go:
--------------------------------------------------------------------------------
1 | // Copyright (c) 2017 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 model
22 |
23 | // CurrentState represents a Helix current state
24 | type CurrentState struct {
25 | ZNRecord
26 | }
27 |
28 | // NewCurrentStateFromMsg creates current state from a message
29 | func NewCurrentStateFromMsg(msg *Message, resourceName string, sessionID string) *CurrentState {
30 | currentState := &CurrentState{*NewRecord(resourceName)}
31 | currentState.SetBucketSize(msg.GetBucketSize())
32 | currentState.SetStateModelDef(msg.GetStateModelDef())
33 | currentState.SetSessionID(sessionID)
34 | currentState.SetBatchMsgMode(msg.GetBatchMsgMode())
35 | currentState.SetStateModelFactoryName(msg.GetStateModelFactoryName())
36 | return currentState
37 | }
38 |
39 | // GetPartitionStateMap returns a map of partition and state
40 | func (s *CurrentState) GetPartitionStateMap() map[string]string {
41 | result := make(map[string]string, len(s.ZNRecord.MapFields))
42 | for partition, val := range s.ZNRecord.MapFields {
43 | result[partition] = val[FieldKeyCurrentState]
44 | }
45 | return result
46 | }
47 |
48 | // GetState returns state of a partition
49 | func (s *CurrentState) GetState(partition string) string {
50 | return s.GetMapField(partition, FieldKeyCurrentState)
51 | }
52 |
53 | // GetSessionID returns the session ID field
54 | func (s *CurrentState) GetSessionID() string {
55 | return s.GetStringField(FieldKeySessionID, "")
56 | }
57 |
58 | // GetStateModelDef return the state model def field
59 | func (s *CurrentState) GetStateModelDef() string {
60 | return s.GetStringField(FieldKeyStateModelDef, "")
61 | }
62 |
63 | // SetBucketSize sets the bucket size field
64 | func (s *CurrentState) SetBucketSize(size int) {
65 | s.SetIntField(FieldKeyBucketSize, size)
66 | }
67 |
68 | // SetSessionID sets the session ID field
69 | func (s *CurrentState) SetSessionID(ID string) {
70 | s.SetSimpleField(FieldKeySessionID, ID)
71 | }
72 |
73 | // SetStateModelDef sets the state model definition field
74 | func (s *CurrentState) SetStateModelDef(stateModelDef string) {
75 | s.SetSimpleField(FieldKeyStateModelDef, stateModelDef)
76 | }
77 |
78 | // SetBatchMsgMode sets the batch message mode field
79 | func (s *CurrentState) SetBatchMsgMode(mode bool) {
80 | s.SetBooleanField(FieldKeyBatchMsgMode, mode)
81 | }
82 |
83 | // SetStateModelFactoryName sets state model factory name field
84 | func (s *CurrentState) SetStateModelFactoryName(name string) {
85 | s.SetSimpleField(FieldKeyStateModelFactoryName, name)
86 | }
87 |
88 | // SetState sets state of a partition
89 | func (s *CurrentState) SetState(partition string, state string) {
90 | s.SetMapField(partition, FieldKeyCurrentState, state)
91 | }
92 |
--------------------------------------------------------------------------------
/model/external_view.go:
--------------------------------------------------------------------------------
1 | // Copyright (c) 2017 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 model
22 |
23 | // ExternalView represents a Helix external view
24 | type ExternalView struct {
25 | ZNRecord
26 | }
27 |
28 | // GetNumPartitions returns the number of partitions for the externalview
29 | func (s *ExternalView) GetNumPartitions() int {
30 | return s.GetIntField(FieldKeyNumPartitions, -1)
31 | }
32 |
--------------------------------------------------------------------------------
/model/ideal_state.go:
--------------------------------------------------------------------------------
1 | // Copyright (c) 2017 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 model
22 |
23 | // IdealState represents a Helix ideal state
24 | type IdealState struct {
25 | ZNRecord
26 | }
27 |
28 | // GetNumPartitions returns the number of partitions for the ideal state
29 | func (s *IdealState) GetNumPartitions() int {
30 | return s.GetIntField(FieldKeyNumPartitions, -1)
31 | }
32 |
--------------------------------------------------------------------------------
/model/instance_config.go:
--------------------------------------------------------------------------------
1 | // Copyright (c) 2017 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 model
22 |
23 | // InstanceConfig represents configs for a Helix instance
24 | type InstanceConfig struct {
25 | ZNRecord
26 | }
27 |
28 | // NewInstanceConfig creates a new instance config property
29 | func NewInstanceConfig(instanceName string) *InstanceConfig {
30 | return &InstanceConfig{*NewRecord(instanceName)}
31 | }
32 |
33 | // SetHost sets host of the instance
34 | func (c *InstanceConfig) SetHost(host string) {
35 | c.SetSimpleField(FieldKeyHelixHost, host)
36 | }
37 |
38 | // SetPort sets port of the instance
39 | func (c *InstanceConfig) SetPort(port int) {
40 | c.SetIntField(FieldKeyHelixPort, port)
41 | }
42 |
43 | // GetEnabled sets if the instance is enabled
44 | func (c *InstanceConfig) GetEnabled() bool {
45 | return c.GetBooleanField(FieldKeyHelixEnabled, false)
46 | }
47 |
48 | // SetEnabled sets if the instance is enabled
49 | func (c *InstanceConfig) SetEnabled(enabled bool) {
50 | c.SetBooleanField(FieldKeyHelixEnabled, enabled)
51 | }
52 |
--------------------------------------------------------------------------------
/model/live_instance.go:
--------------------------------------------------------------------------------
1 | // Copyright (c) 2017 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 model
22 |
23 | import (
24 | "fmt"
25 | "log"
26 | "os"
27 | )
28 |
29 | // LiveInstance represents a Helix participant that is online
30 | type LiveInstance struct {
31 | ZNRecord
32 | }
33 |
34 | // NewLiveInstance creates a new instance of Message for representing a live instance
35 | func NewLiveInstance(instanceName string, sessionID string) *LiveInstance {
36 | liveInstance := &LiveInstance{*NewRecord(instanceName)}
37 | liveInstance.SetSimpleField(FieldKeyHelixVersion, _helixVersion)
38 | liveInstance.SetSimpleField(FieldKeySessionID, sessionID)
39 | liveInstance.SetSimpleField(FieldKeyLiveInstance, getLiveInstanceName(instanceName))
40 | return liveInstance
41 | }
42 |
43 | // GetSessionID returns the session ID field
44 | func (i *LiveInstance) GetSessionID() string {
45 | return i.GetStringField(FieldKeySessionID, "")
46 | }
47 |
48 | // Mirrors Java LiveInstanceName, ManagementFactory.getRuntimeMXBean().getName()
49 | func getLiveInstanceName(instanceName string) string {
50 | hostname, err := os.Hostname()
51 | if err != nil {
52 | log.Printf("failed to get host name for live instance name: %v", err)
53 | hostname = instanceName
54 | }
55 | return fmt.Sprintf("%d@%s", os.Getpid(), hostname)
56 | }
57 |
--------------------------------------------------------------------------------
/model/message.go:
--------------------------------------------------------------------------------
1 | // Copyright (c) 2017 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 model
22 |
23 | import (
24 | "fmt"
25 | "time"
26 |
27 | "github.com/pkg/errors"
28 | )
29 |
30 | // Message helps communication between nodes
31 | // Mirrors org.apache.helix.model.Message
32 | // Sample message:
33 | // {
34 | // "id": "9ff57fc1-9f2a-41a5-af46-c4ae2a54c539",
35 | // "simpleFields": {
36 | // "CREATE_TIMESTAMP": "1425268051457",
37 | // "ClusterEventName": "currentStateChange",
38 | // "FROM_STATE": "OFFLINE",
39 | // "MSG_ID": "9ff57fc1-9f2a-41a5-af46-c4ae2a54c539",
40 | // "MSG_STATE": "new",
41 | // "MSG_TYPE": "STATE_TRANSITION",
42 | // "PARTITION_NAME": "myDB_5",
43 | // "RESOURCE_NAME": "myDB",
44 | // "SRC_NAME": "precise64-CONTROLLER",
45 | // "SRC_SESSION_ID": "14bd852c528004c",
46 | // "STATE_MODEL_DEF": "MasterSlave",
47 | // "STATE_MODEL_FACTORY_NAME": "DEFAULT",
48 | // "TGT_NAME": "localhost_12913",
49 | // "TGT_SESSION_ID": "93406067297878252",
50 | // "TO_STATE": "SLAVE"
51 | // },
52 | // "listFields": {},
53 | // "mapFields": {}
54 | // }
55 | type Message struct {
56 | ZNRecord
57 | }
58 |
59 | // MessageState is the state of the Helix message
60 | // mirrors org.apache.helix.model.Message#MessageState
61 | type MessageState int
62 |
63 | // MessageState is the state of the message
64 | const (
65 | MessageStateNew MessageState = iota
66 | MessageStateRead
67 | MessageStateUnprocessable
68 | )
69 |
70 | // String returns string representation of message state
71 | func (state MessageState) String() string {
72 | switch state {
73 | case MessageStateNew:
74 | return "new"
75 | case MessageStateRead:
76 | return "read"
77 | case MessageStateUnprocessable:
78 | return "unprocessable"
79 | default:
80 | return "unprocessable"
81 | }
82 | }
83 |
84 | // NewMsg creates a new instance of Message instance
85 | func NewMsg(id string) *Message {
86 | return &Message{ZNRecord: *NewRecord(id)}
87 | }
88 |
89 | // SetExecuteStartTime sets execute start time
90 | func (m *Message) SetExecuteStartTime(t time.Time) {
91 | tMs := t.UnixNano() / int64(time.Millisecond)
92 | m.SetSimpleField(FieldKeyExecuteStartTimestamp, fmt.Sprintf("%d", tMs))
93 | }
94 |
95 | // GetStateModelDef returns state model definition
96 | func (m Message) GetStateModelDef() string {
97 | return m.GetStringField(FieldKeyStateModelDef, "")
98 | }
99 |
100 | // SetStateModelDef returns state model definition
101 | func (m *Message) SetStateModelDef(stateModelDef string) {
102 | m.SetSimpleField(FieldKeyStateModelDef, stateModelDef)
103 | }
104 |
105 | // GetMsgState returns the message state
106 | func (m Message) GetMsgState() MessageState {
107 | state := m.GetStringField(FieldKeyMsgState, "")
108 | switch state {
109 | case "new":
110 | return MessageStateNew
111 | case "read":
112 | return MessageStateRead
113 | case "unprocessable":
114 | return MessageStateUnprocessable
115 | default:
116 | return MessageStateUnprocessable
117 | }
118 | }
119 |
120 | // SetMsgState sets the message state
121 | func (m *Message) SetMsgState(s MessageState) {
122 | m.SetSimpleField(FieldKeyMsgState, s.String())
123 | }
124 |
125 | // GetPartitionName returns the partition name
126 | func (m Message) GetPartitionName() (string, error) {
127 | partition := m.GetStringField(FieldKeyPartitionName, "")
128 | if partition == "" {
129 | return "", errors.Errorf("missing partition name in message: %v", m)
130 | }
131 | return partition, nil
132 | }
133 |
134 | // SetPartitionName sets the partition name of the message
135 | func (m *Message) SetPartitionName(partitionName string) {
136 | m.SetSimpleField(FieldKeyPartitionName, partitionName)
137 | }
138 |
139 | // GetTargetSessionID returns target session ID and the message should only be consumed by the
140 | // Helix node of the matching session
141 | func (m Message) GetTargetSessionID() string {
142 | return m.GetStringField(FieldKeyTargetSessionID, "")
143 | }
144 |
145 | // GetTargetName returns the target of the message, such as "PARTICIPANT"
146 | func (m Message) GetTargetName() string {
147 | return m.GetStringField(FieldKeyTargetName, "")
148 | }
149 |
150 | // GetMsgType returns the message type, such as "STATE_TRANSITION"
151 | func (m Message) GetMsgType() string {
152 | return m.GetStringField(FieldKeyMsgType, "")
153 | }
154 |
155 | // GetResourceName returns the resource name
156 | func (m Message) GetResourceName() string {
157 | return m.GetStringField(FieldKeyResourceName, "")
158 | }
159 |
160 | // GetToState returns the toState
161 | func (m Message) GetToState() string {
162 | return m.GetStringField(FieldKeyToState, "")
163 | }
164 |
165 | // GetFromState returns the fromState
166 | func (m Message) GetFromState() string {
167 | return m.GetStringField(FieldKeyFromState, "")
168 | }
169 |
170 | // GetParentMsgID returns the parent message ID
171 | func (m Message) GetParentMsgID() string {
172 | return m.GetStringField(FieldKeyParentMsgID, "")
173 | }
174 |
175 | // GetCreateTimestamp returns the message creation timestamp
176 | func (m Message) GetCreateTimestamp() int64 {
177 | return m.GetInt64Field(FieldKeyCreateTimestamp, 0)
178 | }
179 |
180 | // GetBucketSize return the bucket size of the message
181 | func (m Message) GetBucketSize() int {
182 | return m.GetIntField(FieldKeyBucketSize, 0)
183 | }
184 |
185 | // GetBatchMsgMode returns if the batch message mode should be used
186 | func (m *Message) GetBatchMsgMode() bool {
187 | return m.GetBooleanField(FieldKeyBatchMsgMode, false)
188 | }
189 |
190 | // GetStateModelFactoryName returns the state model factory name
191 | func (m *Message) GetStateModelFactoryName() string {
192 | return m.GetStringField(FieldKeyStateModelFactoryName, _defaultStateModelFactoryName)
193 | }
194 |
--------------------------------------------------------------------------------
/model/model_test.go:
--------------------------------------------------------------------------------
1 | // Copyright (c) 2017 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 model
22 |
23 | import (
24 | "math/rand"
25 | "strconv"
26 | "testing"
27 | "time"
28 |
29 | "github.com/stretchr/testify/assert"
30 | )
31 |
32 | var (
33 | _testMsgJSONString = `{
34 | "id": "9ff57fc1-9f2a-41a5-af46-c4ae2a54c539",
35 | "simpleFields": {
36 | "CREATE_TIMESTAMP": "1425268051457",
37 | "ClusterEventName": "currentStateChange",
38 | "FROM_STATE": "OFFLINE",
39 | "MSG_ID": "9ff57fc1-9f2a-41a5-af46-c4ae2a54c539",
40 | "MSG_STATE": "new",
41 | "MSG_TYPE": "STATE_TRANSITION",
42 | "PARTITION_NAME": "myDB_5",
43 | "RESOURCE_NAME": "myDB",
44 | "SRC_NAME": "precise64-CONTROLLER",
45 | "SRC_SESSION_ID": "14bd852c528004c",
46 | "STATE_MODEL_DEF": "MasterSlave",
47 | "STATE_MODEL_FACTORY_NAME": "DEFAULT",
48 | "TGT_NAME": "localhost_12913",
49 | "TGT_SESSION_ID": "93406067297878252",
50 | "TO_STATE": "SLAVE"
51 | },
52 | "listFields": {},
53 | "mapFields": {}
54 | }`
55 | )
56 |
57 | func TestMsg(t *testing.T) {
58 | msg := NewMsg("test_id")
59 | msg.SetMsgState(MessageStateNew)
60 | assert.Equal(t, MessageStateNew, msg.GetMsgState())
61 | assert.Equal(t, "", msg.GetParentMsgID())
62 | assert.Equal(t, 0, msg.GetBucketSize())
63 | assert.False(t, msg.GetBatchMsgMode())
64 | assert.Equal(t, _defaultStateModelFactoryName, msg.GetStateModelFactoryName())
65 | record, err := NewRecordFromBytes([]byte(_testMsgJSONString))
66 | assert.NoError(t, err)
67 | msg.SetStateModelDef("OnlineOffline")
68 | assert.Equal(t, "OnlineOffline", msg.GetStateModelDef())
69 | msg.SetMsgState(MessageStateRead)
70 | assert.Equal(t, MessageStateRead, msg.GetMsgState())
71 | msg.SetMsgState(MessageStateUnprocessable)
72 | assert.Equal(t, MessageStateUnprocessable, msg.GetMsgState())
73 | msg.SetSimpleField(FieldKeyMsgState, "")
74 | assert.Equal(t, MessageStateUnprocessable, msg.GetMsgState())
75 |
76 | msg = &Message{ZNRecord: *record}
77 | assert.Equal(t, "93406067297878252", msg.GetTargetSessionID())
78 | msg.SetPartitionName("myDB_6")
79 | partition, err := msg.GetPartitionName()
80 | assert.NoError(t, err)
81 | assert.Equal(t, "myDB_6", partition)
82 | assert.Equal(t, "OFFLINE", msg.GetFromState())
83 | assert.Equal(t, "SLAVE", msg.GetToState())
84 | assert.Equal(t, "localhost_12913", msg.GetTargetName())
85 | assert.Equal(t, "STATE_TRANSITION", msg.GetMsgType())
86 | assert.Equal(t, "myDB", msg.GetResourceName())
87 | assert.Equal(t, int64(1425268051457), msg.GetCreateTimestamp())
88 |
89 | record.SetSimpleField(FieldKeyPartitionName, "")
90 | msg = &Message{ZNRecord: *record}
91 | partition, err = msg.GetPartitionName()
92 | assert.Error(t, err)
93 |
94 | now := time.Now()
95 | msg.SetExecuteStartTime(now)
96 | executeStartTime, ok := msg.GetSimpleField(FieldKeyExecuteStartTimestamp)
97 | assert.True(t, ok)
98 | parsed, err := strconv.Atoi(executeStartTime)
99 | assert.NoError(t, err)
100 | assert.Equal(t, now.UnixNano()/int64(time.Millisecond), int64(parsed)) // check rounding
101 | }
102 |
103 | func TestInstanceConfig(t *testing.T) {
104 | config := NewInstanceConfig("test_instance")
105 | assert.False(t, config.GetEnabled())
106 | config.SetEnabled(true)
107 | assert.True(t, config.GetEnabled())
108 | host := "test_host"
109 | port := rand.Int()
110 | config.SetHost(host)
111 | config.SetPort(port)
112 | hostVal, ok := config.GetSimpleField(FieldKeyHelixHost)
113 | assert.True(t, ok)
114 | assert.Equal(t, host, hostVal)
115 | assert.Equal(t, port, config.GetIntField(FieldKeyHelixPort, port+1))
116 | }
117 |
118 | func TestLiveInstanceConfig(t *testing.T) {
119 | instanceName := "test_instance"
120 | instance := NewLiveInstance(instanceName, "test_session")
121 | record := instance.ZNRecord
122 | version, ok := record.GetSimpleField(FieldKeyHelixVersion)
123 | assert.Equal(t, _helixVersion, version)
124 | assert.True(t, ok)
125 | session, ok := record.GetSimpleField(FieldKeySessionID)
126 | assert.Equal(t, "test_session", session)
127 | assert.True(t, ok)
128 | }
129 |
130 | func TestCurrentState(t *testing.T) {
131 | msg := NewMsg("msg_id")
132 | sessionID := "test_session"
133 | state := NewCurrentStateFromMsg(msg, "test_resource", sessionID)
134 | assert.Equal(t, sessionID, state.GetSessionID())
135 |
136 | assert.Equal(t, state.GetState("partition_1"), "")
137 | assert.Equal(t, state.GetState("partition_2"), "")
138 | state.SetState("partition_1", "state1")
139 | state.SetState("partition_2", "state2")
140 | assert.Equal(t, state.GetState("partition_1"), "state1")
141 | assert.Equal(t, state.GetState("partition_2"), "state2")
142 | assert.Len(t, state.GetPartitionStateMap(), 2)
143 | }
144 |
145 | func TestIdealState(t *testing.T) {
146 | numPartitions := 10
147 | record, err := NewRecordFromBytes([]byte("{}"))
148 | assert.NoError(t, err)
149 | record.SetIntField(FieldKeyNumPartitions, numPartitions)
150 | state := &IdealState{ZNRecord: *record}
151 | assert.Equal(t, numPartitions, state.GetNumPartitions())
152 | }
153 |
154 | func TestExternalView(t *testing.T) {
155 | numPartitions := 10
156 | record, err := NewRecordFromBytes([]byte("{}"))
157 | assert.NoError(t, err)
158 | record.SetIntField(FieldKeyNumPartitions, numPartitions)
159 | state := &ExternalView{ZNRecord: *record}
160 | assert.Equal(t, numPartitions, state.GetNumPartitions())
161 | }
162 |
--------------------------------------------------------------------------------
/model/record.go:
--------------------------------------------------------------------------------
1 | // Copyright (c) 2017 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 model
22 |
23 | import (
24 | "encoding/json"
25 | "strconv"
26 | "strings"
27 | )
28 |
29 | const (
30 | _defaultRecordVersion = int32(-1)
31 | )
32 |
33 | // ZNRecord helps communication between nodes
34 | // Mirrors org.apache.helix.ZNRecord
35 | // Sample record:
36 | // {
37 | // "id" : "MyResource",
38 | // "simpleFields" : {
39 | // "REBALANCE_MODE" : "SEMI_AUTO",
40 | // "NUM_PARTITIONS" : "3",
41 | // "REPLICAS" : "2",
42 | // "STATE_MODEL_DEF_REF" : "MasterSlave",
43 | // }
44 | // "listFields" : {
45 | // "MyResource_0" : [node1, node2],
46 | // "MyResource_1" : [node2, node3],
47 | // "MyResource_2" : [node3, node1]
48 | // },
49 | // "mapFields" : {
50 | // }
51 | // }
52 | type ZNRecord struct {
53 | ID string `json:"id"`
54 | SimpleFields map[string]string `json:"simpleFields"`
55 | ListFields map[string][]string `json:"listFields"`
56 | MapFields map[string]map[string]string `json:"mapFields"`
57 | Version int32 `json:"-"`
58 | }
59 |
60 | // NewRecord returns new ZNRecord with id
61 | func NewRecord(id string) *ZNRecord {
62 | return &ZNRecord{
63 | ID: id,
64 | SimpleFields: map[string]string{},
65 | ListFields: map[string][]string{},
66 | MapFields: map[string]map[string]string{},
67 | }
68 | }
69 |
70 | // String returns the beautified JSON string for the ZNRecord
71 | func (r ZNRecord) String() string {
72 | s, _ := r.Marshal()
73 | return string(s)
74 | }
75 |
76 | // Marshal generates the beautified json in byte array format
77 | func (r ZNRecord) Marshal() ([]byte, error) {
78 | return json.MarshalIndent(r, "", " ")
79 | }
80 |
81 | // NewRecordFromBytes creates a new znode instance from a byte array
82 | func NewRecordFromBytes(data []byte) (*ZNRecord, error) {
83 | record := ZNRecord{Version: _defaultRecordVersion}
84 | err := json.Unmarshal(data, &record)
85 | return &record, err
86 | }
87 |
88 | // GetSimpleField returns a value of a key in SimpleField structure
89 | func (r ZNRecord) GetSimpleField(key string) (val string, ok bool) {
90 | if r.SimpleFields == nil {
91 | return "", false
92 | }
93 | val, ok = r.SimpleFields[key]
94 | return val, ok
95 | }
96 |
97 | // GetIntField returns the integer value of a field in the SimpleField
98 | func (r ZNRecord) GetIntField(key string, defaultValue int) int {
99 | value, ok := r.GetSimpleField(key)
100 | if !ok {
101 | return defaultValue
102 | }
103 |
104 | intVal, err := strconv.Atoi(value)
105 | if err != nil {
106 | return defaultValue
107 | }
108 | return intVal
109 | }
110 |
111 | // GetInt64Field returns the integer value of a field in the SimpleField
112 | func (r ZNRecord) GetInt64Field(key string, defaultValue int64) int64 {
113 | value, ok := r.GetSimpleField(key)
114 | if !ok {
115 | return defaultValue
116 | }
117 |
118 | i64, err := strconv.ParseInt(value, 10, 64)
119 | if err != nil {
120 | return defaultValue
121 | }
122 | return i64
123 | }
124 |
125 | // SetIntField sets the integer value of a key under SimpleField.
126 | // the value is stored as in string form
127 | func (r *ZNRecord) SetIntField(key string, value int) {
128 | r.SetSimpleField(key, strconv.Itoa(value))
129 | }
130 |
131 | // GetStringField returns string from simple fields
132 | func (r ZNRecord) GetStringField(key, defaultValue string) string {
133 | res, ok := r.GetSimpleField(key)
134 | if !ok {
135 | return defaultValue
136 | }
137 | return res
138 | }
139 |
140 | // GetBooleanField gets the value of a key under SimpleField and
141 | // convert the result to bool type. That is, if the value is "true",
142 | // the result is true.
143 | func (r ZNRecord) GetBooleanField(key string, defaultValue bool) bool {
144 | result, ok := r.GetSimpleField(key)
145 | if !ok {
146 | return defaultValue
147 | }
148 | return strings.ToLower(result) == "true"
149 | }
150 |
151 | // SetBooleanField sets a key under SimpleField with a specified bool
152 | // value, serialized to string. For example, true will be stored as
153 | // "TRUE"
154 | func (r *ZNRecord) SetBooleanField(key string, value bool) {
155 | r.SetSimpleField(key, strconv.FormatBool(value))
156 | }
157 |
158 | // SetSimpleField sets the value of a key under SimpleField
159 | func (r *ZNRecord) SetSimpleField(key string, value string) {
160 | if r.SimpleFields == nil {
161 | r.SimpleFields = make(map[string]string)
162 | }
163 | r.SimpleFields[key] = value
164 | }
165 |
166 | // RemoveSimpleField removes a key under SimpleField, currently for unit tests only
167 | func (r *ZNRecord) RemoveSimpleField(key string) {
168 | delete(r.SimpleFields, key)
169 | }
170 |
171 | // GetMapField returns the string value of the property of a key
172 | // under MapField.
173 | func (r ZNRecord) GetMapField(key string, property string) string {
174 | if r.MapFields == nil || r.MapFields[key] == nil || r.MapFields[key][property] == "" {
175 | return ""
176 | }
177 | return r.MapFields[key][property]
178 | }
179 |
180 | // SetMapField sets the value of a key under MapField. Both key and
181 | // value are string format.
182 | func (r *ZNRecord) SetMapField(key string, property string, value string) {
183 | if r.MapFields == nil {
184 | r.MapFields = make(map[string]map[string]string)
185 | }
186 |
187 | if r.MapFields[key] == nil {
188 | r.MapFields[key] = make(map[string]string)
189 | }
190 |
191 | r.MapFields[key][property] = value
192 | }
193 |
194 | // RemoveMapField deletes a key from MapField
195 | func (r *ZNRecord) RemoveMapField(key string) {
196 | delete(r.MapFields, key)
197 | }
198 |
--------------------------------------------------------------------------------
/model/record_test.go:
--------------------------------------------------------------------------------
1 | // Copyright (c) 2017 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 model
22 |
23 | import (
24 | "reflect"
25 | "testing"
26 |
27 | "github.com/stretchr/testify/assert"
28 | )
29 |
30 | func TestRecordEncoding(t *testing.T) {
31 | r := &ZNRecord{}
32 | r.Version = _defaultRecordVersion
33 | bytes, err := r.Marshal()
34 | assert.Nil(t, err)
35 | deserialized, err := NewRecordFromBytes(bytes)
36 | assert.True(t, reflect.DeepEqual(r, deserialized))
37 | }
38 |
39 | func TestGetDefaultFields(t *testing.T) {
40 | r := &ZNRecord{}
41 | _, ok := r.GetSimpleField("")
42 | assert.False(t, ok)
43 | defaultInt := 5
44 | i := r.GetIntField("", defaultInt)
45 | assert.Equal(t, i, defaultInt)
46 | r.SetIntField("", defaultInt)
47 | i = r.GetIntField("", defaultInt+1)
48 | assert.Equal(t, i, defaultInt)
49 | boolFieldKey := "bool"
50 | r.SetBooleanField(boolFieldKey, true)
51 | assert.True(t, r.GetBooleanField(boolFieldKey, false))
52 | mapFieldKey, mapFieldProp, mapFieldVal := "k", "p", "val"
53 | r.SetMapField(mapFieldKey, mapFieldProp, mapFieldVal)
54 | assert.Equal(t, mapFieldVal, r.GetMapField(mapFieldKey, mapFieldProp))
55 | }
56 |
--------------------------------------------------------------------------------
/model/state_model_def.go:
--------------------------------------------------------------------------------
1 | // Copyright (c) 2017 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 model
22 |
23 | // StateModelDef represents a Helix ideal state
24 | type StateModelDef struct {
25 | ZNRecord
26 | }
27 |
28 | // GetInitialState returns the initial state for state model def
29 | func (s *StateModelDef) GetInitialState() string {
30 | return s.GetStringField(FieldKeyInitialState, "")
31 | }
32 |
--------------------------------------------------------------------------------
/participant_test.go:
--------------------------------------------------------------------------------
1 | // Copyright (c) 2017 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 helix
22 |
23 | import (
24 | "log"
25 | "math/rand"
26 | "strconv"
27 | "sync"
28 | "testing"
29 | "time"
30 |
31 | "github.com/pkg/errors"
32 | "github.com/samuel/go-zookeeper/zk"
33 | "github.com/stretchr/testify/suite"
34 | "github.com/uber-go/go-helix/model"
35 | "github.com/uber-go/go-helix/util"
36 | uzk "github.com/uber-go/go-helix/zk"
37 | "github.com/uber-go/tally"
38 | "go.uber.org/zap"
39 | )
40 |
41 | type ParticipantTestSuite struct {
42 | BaseHelixTestSuite
43 | }
44 |
45 | func TestParticipantTestSuite(t *testing.T) {
46 | s := &ParticipantTestSuite{}
47 | s.EmbeddedZkPath = "zk/embedded" // TODO (zhijin): remove when moving to github
48 | suite.Run(t, s)
49 | }
50 |
51 | func (s *ParticipantTestSuite) TestConnectAndDisconnect() {
52 | p1, _ := s.createParticipantAndConnect()
53 | err := p1.Connect()
54 | s.NoError(err)
55 | p1.Disconnect()
56 | s.Admin.DropInstance(TestClusterName, p1.instanceName)
57 | s.False(p1.IsConnected())
58 | // disconnect twice doesn't cause panics
59 | p1.Disconnect()
60 |
61 | p2, _ := s.createParticipantAndConnect()
62 | err = p2.createLiveInstance()
63 | s.Error(err, zk.ErrNodeExists)
64 | p2.Disconnect()
65 | s.Admin.DropInstance(TestClusterName, p2.instanceName)
66 | }
67 |
68 | func (s *ParticipantTestSuite) TestParticipantNameCollision() {
69 | port := GetRandomPort()
70 | p1, _ := NewParticipant(zap.NewNop(), tally.NoopScope,
71 | s.ZkConnectString, testApplication, TestClusterName, TestResource, testParticipantHost, port)
72 | s.NotNil(p1)
73 | p1.RegisterStateModel(StateModelNameOnlineOffline, createNoopStateModelProcessor())
74 | err := p1.Connect()
75 | s.NoError(err)
76 | p2, _ := NewParticipant(zap.NewNop(), tally.NoopScope,
77 | s.ZkConnectString, testApplication, TestClusterName, TestResource, testParticipantHost, port)
78 | s.NotNil(p2)
79 | p2.RegisterStateModel(StateModelNameOnlineOffline, createNoopStateModelProcessor())
80 | err = p2.Connect()
81 | s.Error(err)
82 | }
83 |
84 | func (s *ParticipantTestSuite) TestIsClusterSetup() {
85 | p, _ := s.createParticipantAndConnect()
86 | defer p.Disconnect()
87 |
88 | setup, err := p.isClusterSetup()
89 | s.NoError(err)
90 | s.True(setup)
91 | }
92 |
93 | func (s *ParticipantTestSuite) TestHandleInvalidMessages() {
94 | p, _ := s.createParticipantAndConnect()
95 | defer p.Disconnect()
96 |
97 | msgEmpty := &model.Message{}
98 | err := p.handleMsg(msgEmpty)
99 | s.Equal(errMsgMissingStateModelDef, err)
100 |
101 | msgWithoutStateModelDef := s.createMsg(p, removeMsgFieldsOp(model.FieldKeyStateModelDef))
102 | err = p.handleMsg(msgWithoutStateModelDef)
103 | s.Equal(errMsgMissingStateModelDef, err)
104 |
105 | msgWithoutFromState := s.createMsg(p, removeMsgFieldsOp(model.FieldKeyFromState))
106 | err = p.handleMsg(msgWithoutFromState)
107 | s.Equal(errMsgMissingFromOrToState, err)
108 |
109 | msgWithoutToState := s.createMsg(p, removeMsgFieldsOp(model.FieldKeyToState))
110 | err = p.handleMsg(msgWithoutToState)
111 | s.Equal(errMsgMissingFromOrToState, err)
112 | }
113 |
114 | func (s *ParticipantTestSuite) TestHandleValidMessages() {
115 | p, _ := s.createParticipantAndConnect()
116 | defer p.Disconnect()
117 |
118 | validMsgs := []*model.Message{
119 | s.createMsg(p),
120 | s.createMsg(p, removeMsgFieldsOp(model.FieldKeyParentMsgID)),
121 | s.createMsg(p, setMsgFieldsOp(model.FieldKeyTargetSessionID, "wrong id")),
122 | s.createMsg(p, setMsgFieldsOp(model.FieldKeyToState, StateModelStateDropped)),
123 | }
124 |
125 | for _, msg := range validMsgs {
126 | err := p.handleMsg(msg)
127 | s.NoError(err)
128 | }
129 | }
130 |
131 | func (s *ParticipantTestSuite) TestGetCurrentResourceNames() {
132 | p, _ := s.createParticipantAndConnect()
133 | defer p.Disconnect()
134 |
135 | resources := util.NewStringSet(p.getCurrentResourceNames()...)
136 | s.Equal(0, resources.Size())
137 |
138 | keyBuilder := &KeyBuilder{TestClusterName}
139 | accessor := p.DataAccessor()
140 |
141 | resource := CreateRandomString()
142 | msg := model.NewMsg("test_id")
143 | session := p.zkClient.GetSessionID()
144 | path := keyBuilder.currentStateForResource(p.instanceName, session, resource)
145 | state := model.NewCurrentStateFromMsg(msg, resource, session)
146 | err := accessor.createCurrentState(path, state)
147 | s.NoError(err)
148 |
149 | resources = util.NewStringSet(p.getCurrentResourceNames()...)
150 | s.Equal(1, resources.Size(), "participant should find one resource")
151 | s.True(resources.Contains(resource), "participant should find the added resource")
152 | }
153 |
154 | func (s *ParticipantTestSuite) TestCarryOverPreviousCurrentStateWhenCurrentStateNotExist() {
155 | p, _ := s.createParticipantAndConnect()
156 | defer p.Disconnect()
157 |
158 | keyBuilder := &KeyBuilder{TestClusterName}
159 | client := s.CreateAndConnectClient()
160 | defer client.Disconnect()
161 | accessor := newDataAccessor(client, keyBuilder)
162 | resource := CreateRandomString()
163 | preSession := CreateRandomString()
164 | currentSession := p.zkClient.GetSessionID()
165 |
166 | // Current state is not created
167 | _, err := accessor.CurrentState(p.instanceName, currentSession, resource)
168 | s.EqualError(errors.Cause(err), zk.ErrNoNode.Error())
169 |
170 | path := keyBuilder.currentStateForResource(p.instanceName, preSession, resource)
171 | preState := &model.CurrentState{ZNRecord: *model.NewRecord(resource)}
172 | preState.SetStateModelDef(StateModelNameOnlineOffline)
173 | preState.SetState("partition_1", StateModelStateOffline)
174 | preState.SetState("partition_2", StateModelStateOnline)
175 | preState.SetState("partition_3", StateModelStateOnline)
176 | err = accessor.createCurrentState(path, preState)
177 | s.NoError(err)
178 |
179 | p.carryOverPreviousCurrentState()
180 | currentState, err := accessor.CurrentState(p.instanceName, currentSession, resource)
181 | s.NoError(err)
182 | // carry over only when current state not exist
183 | // if not exist, set to initial state
184 | s.Equal(currentState.GetState("partition_1"), StateModelStateOffline)
185 | s.Equal(currentState.GetState("partition_2"), StateModelStateOffline)
186 | s.Equal(currentState.GetState("partition_3"), StateModelStateOffline)
187 | s.Len(currentState.GetPartitionStateMap(), 3)
188 | }
189 |
190 | func (s *ParticipantTestSuite) TestCarryOverPreviousCurrentStateWhenCurrentStateExists() {
191 | p, _ := s.createParticipantAndConnect()
192 | defer p.Disconnect()
193 |
194 | keyBuilder := &KeyBuilder{TestClusterName}
195 | client := s.CreateAndConnectClient()
196 | defer client.Disconnect()
197 | accessor := newDataAccessor(client, keyBuilder)
198 | resource := CreateRandomString()
199 | preSession := CreateRandomString()
200 | currentSession := p.zkClient.GetSessionID()
201 |
202 | // Current state is not created
203 | _, err := accessor.CurrentState(p.instanceName, currentSession, resource)
204 | s.EqualError(errors.Cause(err), zk.ErrNoNode.Error())
205 |
206 | path := keyBuilder.currentStateForResource(p.instanceName, preSession, resource)
207 | preState := &model.CurrentState{ZNRecord: *model.NewRecord(resource)}
208 | preState.SetStateModelDef(StateModelNameOnlineOffline)
209 | preState.SetState("partition_1", StateModelStateOffline)
210 | preState.SetState("partition_2", StateModelStateOnline)
211 | preState.SetState("partition_3", StateModelStateOnline)
212 | err = accessor.createCurrentState(path, preState)
213 | s.NoError(err)
214 |
215 | path = keyBuilder.currentStateForResource(p.instanceName, currentSession, resource)
216 | currentState := &model.CurrentState{ZNRecord: *model.NewRecord(resource)}
217 | currentState.SetStateModelDef(StateModelNameOnlineOffline)
218 | currentState.SetState("partition_2", StateModelStateOnline)
219 | err = accessor.createCurrentState(path, currentState)
220 | s.NoError(err)
221 |
222 | p.carryOverPreviousCurrentState()
223 | currentState, err = accessor.CurrentState(p.instanceName, currentSession, resource)
224 | s.NoError(err)
225 | // carry over only when current state not exist
226 | // if not exist, set to initial state
227 | s.Equal(currentState.GetState("partition_1"), StateModelStateOffline)
228 | s.Equal(currentState.GetState("partition_2"), StateModelStateOnline)
229 | s.Equal(currentState.GetState("partition_3"), StateModelStateOffline)
230 | s.Len(currentState.GetPartitionStateMap(), 3)
231 | }
232 |
233 | func (s *ParticipantTestSuite) TestCarryOverPreviousCurrentStateForMultipleResources() {
234 | p, _ := s.createParticipantAndConnect()
235 | defer p.Disconnect()
236 |
237 | numberOfResources := 3
238 | var resources []string
239 | keyBuilder := &KeyBuilder{TestClusterName}
240 | client := s.CreateAndConnectClient()
241 | defer client.Disconnect()
242 | accessor := newDataAccessor(client, keyBuilder)
243 | preSession := CreateRandomString()
244 | currentSession := p.zkClient.GetSessionID()
245 |
246 | for i := 0; i < numberOfResources; i++ {
247 | resource := CreateRandomString()
248 | resources = append(resources, resource)
249 | path := keyBuilder.currentStateForResource(p.instanceName, preSession, resource)
250 | preState := &model.CurrentState{ZNRecord: *model.NewRecord(resource)}
251 | preState.SetStateModelDef(StateModelNameOnlineOffline)
252 | preState.SetState("partition_1", StateModelStateOffline)
253 | preState.SetState("partition_2", StateModelStateOnline)
254 | preState.SetState("partition_3", StateModelStateOnline)
255 | err := accessor.createCurrentState(path, preState)
256 | s.NoError(err)
257 | }
258 |
259 | p.carryOverPreviousCurrentState()
260 |
261 | for _, resource := range resources {
262 | currentState, err := accessor.CurrentState(p.instanceName, currentSession, resource)
263 | s.NoError(err)
264 | // carry over only when current state not exist
265 | // if not exist, set to initial state
266 | s.Equal(currentState.GetState("partition_1"), StateModelStateOffline)
267 | s.Equal(currentState.GetState("partition_2"), StateModelStateOffline)
268 | s.Equal(currentState.GetState("partition_3"), StateModelStateOffline)
269 | s.Len(currentState.GetPartitionStateMap(), 3)
270 | }
271 | }
272 |
273 | func (s *ParticipantTestSuite) TestProcessMessages() {
274 | port := GetRandomPort()
275 | processor := NewStateModelProcessor()
276 | mu := &sync.Mutex{}
277 | counters := map[string]map[string]int{
278 | StateModelStateOffline: {
279 | StateModelStateOnline: 0,
280 | StateModelStateDropped: 0,
281 | },
282 | StateModelStateOnline: {
283 | StateModelStateOffline: 0,
284 | },
285 | }
286 |
287 | processor.AddTransition(
288 | StateModelStateOnline, StateModelStateOffline, func(m *model.Message) error {
289 | mu.Lock()
290 | defer mu.Unlock()
291 | counters[StateModelStateOnline][StateModelStateOffline]++
292 | log.Printf("partition ONLINE=>OFFLINE: %v", m)
293 | return nil
294 | })
295 | processor.AddTransition(
296 | StateModelStateOffline, StateModelStateOnline, func(m *model.Message) error {
297 | mu.Lock()
298 | defer mu.Unlock()
299 | counters[StateModelStateOffline][StateModelStateOnline]++
300 | log.Printf("partition OFFLINE=>ONLINE: %v", m)
301 | return nil
302 | })
303 | processor.AddTransition(
304 | StateModelStateOffline, StateModelStateDropped, func(m *model.Message) error {
305 | mu.Lock()
306 | defer mu.Unlock()
307 | counters[StateModelStateOffline][StateModelStateDropped]++
308 | return nil
309 | })
310 | p, _ := NewParticipant(zap.NewNop(), tally.NoopScope,
311 | s.ZkConnectString, testApplication, TestClusterName, TestResource, testParticipantHost, port)
312 | pImpl := p.(*participant)
313 | s.NotNil(pImpl)
314 | pImpl.RegisterStateModel(StateModelNameOnlineOffline, processor)
315 | err := pImpl.Connect()
316 | s.NoError(err)
317 | mu.Lock()
318 | s.Equal(0, counters[StateModelStateOnline][StateModelStateOffline])
319 | s.Equal(0, counters[StateModelStateOffline][StateModelStateOnline])
320 | s.Equal(0, counters[StateModelStateOffline][StateModelStateDropped])
321 | mu.Unlock()
322 |
323 | keyBuilder := &KeyBuilder{TestClusterName}
324 | client := s.CreateAndConnectClient()
325 | defer client.Disconnect()
326 | accessor := newDataAccessor(client, keyBuilder)
327 |
328 | msg := s.createMsg(pImpl,
329 | setMsgFieldsOp(model.FieldKeyFromState, StateModelStateOffline),
330 | setMsgFieldsOp(model.FieldKeyToState, StateModelStateOnline),
331 | )
332 | accessor.createMsg(keyBuilder.participantMsg(pImpl.InstanceName(), CreateRandomString()), msg)
333 | // wait for the participant to process messages
334 | time.Sleep(1 * time.Second)
335 | mu.Lock()
336 | s.Equal(1, counters[StateModelStateOffline][StateModelStateOnline])
337 | s.Equal(0, counters[StateModelStateOffline][StateModelStateDropped])
338 | s.Equal(0, counters[StateModelStateOnline][StateModelStateOffline])
339 | mu.Unlock()
340 |
341 | msgID := CreateRandomString()
342 | msg = s.createMsg(pImpl,
343 | setMsgFieldsOp(model.FieldKeyTargetSessionID, CreateRandomString()),
344 | )
345 | msg.ID = msgID
346 | accessor.createMsg(keyBuilder.participantMsg(pImpl.InstanceName(), msgID), msg)
347 | time.Sleep(2 * time.Second)
348 | path := keyBuilder.participantMsg(pImpl.InstanceName(), msgID)
349 | exists, _, err := client.Exists(path)
350 | s.NoError(err)
351 | s.False(exists)
352 |
353 | msgID = CreateRandomString()
354 | msg = s.createMsg(pImpl,
355 | setMsgFieldsOp(model.FieldKeyMsgType, MsgTypeNoop),
356 | )
357 | msg.ID = msgID
358 | accessor.createMsg(keyBuilder.participantMsg(pImpl.InstanceName(), msgID), msg)
359 | time.Sleep(2 * time.Second)
360 | path = keyBuilder.participantMsg(pImpl.InstanceName(), msgID)
361 | exists, _, err = client.Exists(path)
362 | s.NoError(err)
363 | s.False(exists)
364 | }
365 |
366 | func (s *ParticipantTestSuite) TestUpdateCurrentState() {
367 | p, _ := s.createParticipantAndConnect()
368 | defer p.Disconnect()
369 |
370 | keyBuilder := &KeyBuilder{TestClusterName}
371 | client := s.CreateAndConnectClient()
372 | defer client.Disconnect()
373 | accessor := newDataAccessor(client, keyBuilder)
374 |
375 | currentStatePath := keyBuilder.currentStatesForSession(p.instanceName, p.zkClient.GetSessionID())
376 | exists, _, err := client.Exists(currentStatePath)
377 | s.NoError(err)
378 | s.False(exists)
379 |
380 | resource := CreateRandomString()
381 | partition := strconv.Itoa(rand.Int())
382 | msg := s.createMsg(p,
383 | setMsgFieldsOp(model.FieldKeyFromState, StateModelStateOffline),
384 | setMsgFieldsOp(model.FieldKeyToState, StateModelStateOnline),
385 | setMsgFieldsOp(model.FieldKeyResourceName, resource),
386 | setMsgFieldsOp(model.FieldKeyMsgType, MsgTypeStateTransition),
387 | setMsgFieldsOp(model.FieldKeyPartitionName, partition),
388 | )
389 | accessor.CreateParticipantMsg(p.instanceName, msg)
390 | // wait for the participant to process messages
391 | time.Sleep(2 * time.Second)
392 |
393 | currentState, err := accessor.CurrentState(p.instanceName, p.zkClient.GetSessionID(), resource)
394 | s.NoError(err)
395 | s.Equal(int32(1), currentState.Version, "current state has been created and updated once")
396 | s.Equal(currentState.GetState(partition), StateModelStateOnline)
397 | }
398 |
399 | func (s *ParticipantTestSuite) TestMismatchStateIsRejected() {
400 | p, _ := s.createParticipantAndConnect()
401 | defer p.Disconnect()
402 |
403 | keyBuilder := &KeyBuilder{TestClusterName}
404 | client := s.CreateAndConnectClient()
405 | defer client.Disconnect()
406 | accessor := newDataAccessor(client, keyBuilder)
407 |
408 | currentStatePath := keyBuilder.currentStatesForSession(p.instanceName, p.zkClient.GetSessionID())
409 | exists, _, err := client.Exists(currentStatePath)
410 | s.NoError(err)
411 | s.False(exists)
412 |
413 | resource := CreateRandomString()
414 | partition := strconv.Itoa(rand.Int())
415 | msg := s.createMsg(p,
416 | setMsgFieldsOp(model.FieldKeyFromState, StateModelStateOffline),
417 | setMsgFieldsOp(model.FieldKeyToState, StateModelStateOnline),
418 | setMsgFieldsOp(model.FieldKeyResourceName, resource),
419 | setMsgFieldsOp(model.FieldKeyMsgType, MsgTypeStateTransition),
420 | setMsgFieldsOp(model.FieldKeyPartitionName, partition),
421 | )
422 | accessor.createMsg(keyBuilder.participantMsg(p.instanceName, CreateRandomString()), msg)
423 | // wait for the participant to process messages
424 | time.Sleep(2 * time.Second)
425 | currentState, err := accessor.CurrentState(p.instanceName, p.zkClient.GetSessionID(), resource)
426 | s.NoError(err)
427 | s.Equal(currentState.GetState(partition), StateModelStateOnline)
428 |
429 | // the expected fromState is ONLINE, so the transition should be rejected
430 | msg = s.createMsg(p,
431 | setMsgFieldsOp(model.FieldKeyFromState, StateModelStateOffline),
432 | setMsgFieldsOp(model.FieldKeyToState, StateModelStateDropped),
433 | setMsgFieldsOp(model.FieldKeyResourceName, resource),
434 | setMsgFieldsOp(model.FieldKeyMsgType, MsgTypeStateTransition),
435 | setMsgFieldsOp(model.FieldKeyPartitionName, partition),
436 | )
437 | accessor.createMsg(keyBuilder.participantMsg(p.instanceName, CreateRandomString()), msg)
438 | // wait for the participant to process messages
439 | time.Sleep(2 * time.Second)
440 | currentState, err = accessor.CurrentState(p.instanceName, p.zkClient.GetSessionID(), resource)
441 | s.NoError(err)
442 | // state should remain unchanged
443 | s.Equal(currentState.GetState(partition), StateModelStateOnline)
444 | }
445 |
446 | func (s *ParticipantTestSuite) TestHandleNewSessionCalledAfterZookeeperSessionExpired() {
447 | port := GetRandomPort()
448 | p, _ := NewParticipant(zap.NewNop(), tally.NoopScope,
449 | s.ZkConnectString, testApplication, TestClusterName, TestResource, testParticipantHost, port)
450 | pImpl := p.(*participant)
451 | defer pImpl.Disconnect()
452 | // set default state to StateHasSession so Participant would believe it is connected
453 | fakeZK := uzk.NewFakeZk(uzk.DefaultConnectionState(zk.StateHasSession))
454 | pImpl.zkClient = uzk.NewClient(zap.NewNop(), tally.NoopScope, uzk.WithConnFactory(fakeZK),
455 | uzk.WithRetryTimeout(time.Second))
456 | pImpl.Connect()
457 | s.Len(fakeZK.GetConnections(), 1)
458 |
459 | // Create(liveInstancePath) will be called if and only if a new session is created
460 | // check the times it is called would infer the number of times a new session is created
461 | fakeZKConnection := fakeZK.GetConnections()[0]
462 | liveInstancePath := pImpl.keyBuilder.liveInstance(pImpl.InstanceName())
463 | methodHistory := fakeZKConnection.GetHistory()
464 | historyForChildrenW := methodHistory.GetHistoryForMethod("Create")
465 | s.Len(historyForChildrenW, 1)
466 | s.Equal(historyForChildrenW[0].Params[0].(string), liveInstancePath)
467 |
468 | // simulate a session expiration and reconnection event
469 | fakeZK.SetState(fakeZKConnection, zk.StateExpired)
470 | fakeZK.SetState(fakeZKConnection, zk.StateHasSession)
471 | time.Sleep(1 * time.Second)
472 | historyForChildrenW = methodHistory.GetHistoryForMethod("Create")
473 | s.Len(historyForChildrenW, 2)
474 | s.Equal(historyForChildrenW[1].Params[0].(string), liveInstancePath)
475 | }
476 |
477 | func (s *ParticipantTestSuite) TestFatalErrorCh() {
478 | port := GetRandomPort()
479 | p, errCh := NewTestParticipant(zap.NewNop(), tally.NoopScope,
480 | s.ZkConnectString, testApplication, TestClusterName, TestResource, testParticipantHost, port)
481 | s.True(errCh == p.GetFatalErrorChan())
482 | }
483 |
484 | func (s *ParticipantTestSuite) createMsg(p *participant, ops ...msgOp) *model.Message {
485 | msg := s.createValidMsg(p)
486 | for _, op := range ops {
487 | op(msg)
488 | }
489 | // TODO: add msg to zk?
490 | return msg
491 | }
492 |
493 | func (s *ParticipantTestSuite) createValidMsg(p *participant) *model.Message {
494 | msg := model.NewMsg(CreateRandomString())
495 | msg.SetSimpleField(model.FieldKeyStateModelDef, StateModelNameOnlineOffline)
496 | msg.SetSimpleField(model.FieldKeyTargetSessionID, p.zkClient.GetSessionID())
497 | msg.SetSimpleField(model.FieldKeyFromState, StateModelStateOffline)
498 | msg.SetSimpleField(model.FieldKeyToState, StateModelStateOnline)
499 | msg.SetMsgState(model.MessageStateNew)
500 | return msg
501 | }
502 |
503 | type msgOp func(*model.Message)
504 |
505 | func removeMsgFieldsOp(fields ...string) msgOp {
506 | return func(msg *model.Message) {
507 | for _, f := range fields {
508 | msg.RemoveSimpleField(f)
509 | msg.RemoveMapField(f)
510 | }
511 | }
512 | }
513 |
514 | func setMsgFieldsOp(key, val string) msgOp {
515 | return func(msg *model.Message) {
516 | msg.SetSimpleField(key, val)
517 | }
518 | }
519 |
--------------------------------------------------------------------------------
/state_model.go:
--------------------------------------------------------------------------------
1 | // Copyright (c) 2017 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 helix
22 |
23 | import "sync"
24 |
25 | // StateModel mirrors the behavior of org.apache.helix.participant.statemachine.StateModel and
26 | // org.apache.helix.participant.statemachine.StateModelFactory
27 | // it keeps the state of each resource/partition combination
28 | type StateModel interface {
29 | // GetState returns the state of a resource/partition combination
30 | GetState(resourceName string, partitionKey string) (string, bool)
31 | // UpdateState updates the state of a resource/partition combination
32 | UpdateState(resourceName string, partitionKey string, state string)
33 | // RemoveState removes the state of a resource/partition combination
34 | RemoveState(resourceName string, partitionKey string)
35 | }
36 |
37 | type stateModel struct {
38 | sync.RWMutex
39 | stateModelMap map[string]map[string]string
40 | }
41 |
42 | // NewStateModel creates a StateModel
43 | func NewStateModel() StateModel {
44 | return &stateModel{
45 | stateModelMap: make(map[string]map[string]string),
46 | }
47 | }
48 |
49 | func (s *stateModel) GetState(resourceName string, partitionKey string) (string, bool) {
50 | s.RLock()
51 | defer s.RUnlock()
52 | partitionStateMap, exists := s.stateModelMap[resourceName]
53 | if !exists {
54 | return "", false
55 | }
56 | state, exists := partitionStateMap[partitionKey]
57 | if !exists {
58 | return "", false
59 | }
60 | return state, true
61 | }
62 |
63 | func (s *stateModel) UpdateState(resourceName string, partitionKey string, state string) {
64 | s.Lock()
65 | defer s.Unlock()
66 | partitionStateMap, exists := s.stateModelMap[resourceName]
67 | if !exists {
68 | partitionStateMap = make(map[string]string)
69 | s.stateModelMap[resourceName] = partitionStateMap
70 | }
71 | partitionStateMap[partitionKey] = state
72 | }
73 |
74 | func (s *stateModel) RemoveState(resourceName string, partitionKey string) {
75 | s.Lock()
76 | defer s.Unlock()
77 | partitionStateMap, exists := s.stateModelMap[resourceName]
78 | if !exists {
79 | return
80 | }
81 | delete(partitionStateMap, partitionKey)
82 | if len(partitionStateMap) == 0 {
83 | delete(s.stateModelMap, resourceName)
84 | }
85 | }
86 |
--------------------------------------------------------------------------------
/state_model_processor.go:
--------------------------------------------------------------------------------
1 | // Copyright (c) 2017 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 helix
22 |
23 | import (
24 | "github.com/uber-go/go-helix/model"
25 | )
26 |
27 | // StateTransitionHandler is type for handler method
28 | type StateTransitionHandler func(msg *model.Message) error
29 |
30 | // Transition associates a handler function with state transitions
31 | type Transition struct {
32 | FromState string
33 | ToState string
34 | Handler StateTransitionHandler
35 | }
36 |
37 | // StateModelProcessor handles state transitions
38 | // This mirrors org.apache.helix.participant.statemachine.StateModelFactory
39 | type StateModelProcessor struct {
40 | // fromState->toState->StateTransitionHandler
41 | Transitions map[string]map[string]StateTransitionHandler
42 | }
43 |
44 | // NewStateModelProcessor functions similarly to StateMachineEngine
45 | func NewStateModelProcessor() *StateModelProcessor {
46 | return &StateModelProcessor{
47 | Transitions: map[string]map[string]StateTransitionHandler{},
48 | }
49 | }
50 |
51 | // AddTransition adds a new transition handler
52 | func (p *StateModelProcessor) AddTransition(fromState string, toState string, handler StateTransitionHandler) {
53 | if _, ok := p.Transitions[fromState]; !ok {
54 | p.Transitions[fromState] = make(map[string]StateTransitionHandler)
55 | }
56 | p.Transitions[fromState][toState] = handler
57 | }
58 |
--------------------------------------------------------------------------------
/state_model_test.go:
--------------------------------------------------------------------------------
1 | // Copyright (c) 2017 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 helix
22 |
23 | import (
24 | "testing"
25 |
26 | "github.com/stretchr/testify/suite"
27 | )
28 |
29 | const (
30 | _testResource string = "testResource"
31 | _testPartition string = "testPartition"
32 | _testPartitionTwo string = "testPartitionTwo"
33 | )
34 |
35 | type StateModelTestSuite struct {
36 | suite.Suite
37 | }
38 |
39 | func TestStateModelTestSuite(t *testing.T) {
40 | suite.Run(t, &StateModelTestSuite{})
41 | }
42 |
43 | func (s *StateModelTestSuite) TestUpdateAndGetState() {
44 | stateModel := NewStateModel()
45 | result, exist := stateModel.GetState(_testResource, _testPartition)
46 | s.False(exist)
47 |
48 | stateModel.UpdateState(_testResource, _testPartition, StateModelStateOffline)
49 | result, exist = stateModel.GetState(_testResource, _testPartition)
50 | s.True(exist)
51 | s.Equal(result, StateModelStateOffline)
52 |
53 | stateModel.UpdateState(_testResource, _testPartition, StateModelStateOnline)
54 | result, exist = stateModel.GetState(_testResource, _testPartition)
55 | s.True(exist)
56 | s.Equal(result, StateModelStateOnline)
57 | }
58 |
59 | func (s *StateModelTestSuite) TestRemoveState() {
60 | stateModel := NewStateModel()
61 | result, exist := stateModel.GetState(_testResource, _testPartition)
62 | s.False(exist)
63 |
64 | stateModel.UpdateState(_testResource, _testPartition, StateModelStateOffline)
65 | result, exist = stateModel.GetState(_testResource, _testPartition)
66 | s.True(exist)
67 | s.Equal(result, StateModelStateOffline)
68 |
69 | stateModel.RemoveState(_testResource, _testPartition)
70 | _, exist = stateModel.GetState(_testResource, _testPartition)
71 | s.False(exist)
72 | }
73 |
74 | func (s *StateModelTestSuite) TestRemoveStateWillCleanSubMap() {
75 | stateModel := NewStateModel().(*stateModel)
76 | stateModel.UpdateState(_testResource, _testPartition, StateModelStateOffline)
77 | stateModel.UpdateState(_testResource, _testPartitionTwo, StateModelStateOffline)
78 | s.Len(stateModel.stateModelMap, 1)
79 |
80 | stateModel.RemoveState(_testResource, _testPartition)
81 | stateModel.RemoveState(_testResource, _testPartitionTwo)
82 | s.Len(stateModel.stateModelMap, 0)
83 | }
84 |
85 | func (s *StateModelTestSuite) TestRemoveNonExistentState() {
86 | stateModel := NewStateModel()
87 | stateModel.RemoveState(_testResource, _testPartition)
88 | _, exist := stateModel.GetState(_testResource, _testPartition)
89 | s.False(exist)
90 |
91 | stateModel.RemoveState(_testResource, _testPartition)
92 | _, exist = stateModel.GetState(_testResource, _testPartition)
93 | s.False(exist)
94 | }
95 |
--------------------------------------------------------------------------------
/test_participant.go:
--------------------------------------------------------------------------------
1 | // Copyright (c) 2017 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 helix
22 |
23 | import (
24 | "github.com/uber-go/tally"
25 | "go.uber.org/zap"
26 | )
27 |
28 | // TestParticipant is a participant used for test purpose
29 | type TestParticipant struct {
30 | Participant
31 | }
32 |
33 | // NewTestParticipant returns a TestParticipant
34 | func NewTestParticipant(
35 | logger *zap.Logger,
36 | scope tally.Scope,
37 | zkConnectString string,
38 | application string,
39 | clusterName string,
40 | resourceName string,
41 | host string,
42 | port int32,
43 | ) (*TestParticipant, <-chan error) {
44 | participant, fatalErrChan :=
45 | NewParticipant(logger, scope, zkConnectString, application, clusterName, resourceName, host, port)
46 | return &TestParticipant{participant}, fatalErrChan
47 | }
48 |
49 | // GetFatalErrorChan returns the fatal error chan, so user can send test error
50 | func (p *TestParticipant) GetFatalErrorChan() chan error {
51 | return p.Participant.(*participant).fatalErrChan
52 | }
53 |
--------------------------------------------------------------------------------
/test_util.go:
--------------------------------------------------------------------------------
1 | // Copyright (c) 2017 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 helix
22 |
23 | import (
24 | "fmt"
25 | "log"
26 | "math/rand"
27 |
28 | "github.com/uber-go/go-helix/model"
29 | "github.com/uber-go/go-helix/zk"
30 | "github.com/uber-go/tally"
31 | "go.uber.org/zap"
32 | )
33 |
34 | const (
35 | // TestClusterName is the name of cluster used in test
36 | TestClusterName = "test_cluster"
37 | // TestResource is the name of resource used in test
38 | TestResource = "test_resource"
39 | testApplication = "test_app"
40 | testParticipantHost = "localhost"
41 | testInstanceName = "test_instance"
42 | maxValidPort = int32(2<<15) - 1
43 | )
44 |
45 | // BaseHelixTestSuite can be embedded in any test suites that need to
46 | // interact with the test zk server or helix admin. It provides ZK
47 | // connect string and Helix admin references for convenient uses in parent suites.
48 | type BaseHelixTestSuite struct {
49 | zk.BaseZkTestSuite
50 | Admin *Admin
51 | }
52 |
53 | // SetupSuite ensures zk server is up
54 | func (s *BaseHelixTestSuite) SetupSuite() {
55 | s.BaseZkTestSuite.SetupSuite()
56 | admin, err := NewAdmin(s.ZkConnectString)
57 | s.Admin = admin
58 | s.NoError(err)
59 |
60 | s.ensureHelixClusterUp()
61 | }
62 |
63 | // TearDownSuite disconnects zk if not done already
64 | func (s *BaseHelixTestSuite) TearDownSuite() {
65 | if s.Admin.zkClient.IsConnected() {
66 | s.Admin.zkClient.Disconnect()
67 | }
68 | }
69 |
70 | // GetRandomPort returns random valid port number (1~65535)
71 | func GetRandomPort() int32 {
72 | return rand.Int31n(maxValidPort) + 1
73 | }
74 |
75 | func (s *BaseHelixTestSuite) ensureHelixClusterUp() {
76 | setup, err := s.Admin.isClusterSetup(TestClusterName)
77 | s.NoError(err, "Error checking helix cluster setup")
78 | if !setup {
79 | log.Println("test helix cluster doesn't exist, try creating..")
80 | created := s.Admin.AddCluster(TestClusterName, true)
81 | if !created {
82 | s.Fail("Error creating helix cluster")
83 | }
84 | }
85 | err = s.Admin.SetConfig(TestClusterName, "CLUSTER", map[string]string{
86 | _allowParticipantAutoJoinKey: "true",
87 | })
88 | s.NoError(err, "error setting cluster config")
89 | }
90 |
91 | func (s *BaseHelixTestSuite) createParticipantAndConnect() (*participant, <-chan error) {
92 | port := GetRandomPort()
93 | p, errChan := NewParticipant(zap.NewNop(), tally.NoopScope,
94 | s.ZkConnectString, testApplication, TestClusterName, TestResource, testParticipantHost, port)
95 | pImpl := p.(*participant)
96 | s.NotNil(p)
97 | p.RegisterStateModel(StateModelNameOnlineOffline, createNoopStateModelProcessor())
98 | log.Println("Created participant ", pImpl.instanceName)
99 | err := p.Connect()
100 | s.NoError(err)
101 | return pImpl, errChan
102 | }
103 |
104 | func createNoopStateModelProcessor() *StateModelProcessor {
105 | processor := NewStateModelProcessor()
106 | processor.AddTransition(
107 | StateModelStateOnline, StateModelStateOffline, func(m *model.Message) error {
108 | log.Printf("partition ONLINE=>OFFLINE: %v", m)
109 | return nil
110 | })
111 | processor.AddTransition(
112 | StateModelStateOffline, StateModelStateOnline, func(m *model.Message) error {
113 | log.Printf("partition OFFLINE=>ONLINE: %v", m)
114 | return nil
115 | })
116 | processor.AddTransition(
117 | StateModelStateOffline, StateModelStateDropped, func(m *model.Message) error {
118 | log.Printf("partition OFFLINE=>DROPPED: %v", m)
119 | return nil
120 | })
121 | return processor
122 | }
123 |
124 | // CreateRandomString creates a random with numeric characters
125 | func CreateRandomString() string {
126 | return fmt.Sprintf("%d", rand.Int63())
127 | }
128 |
--------------------------------------------------------------------------------
/util/set.go:
--------------------------------------------------------------------------------
1 | // Copyright (c) 2017 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 util
22 |
23 | // StringSet is a wrapper for a hash set that stores string keys
24 | type StringSet map[string]struct{}
25 |
26 | // NewStringSet create a new set
27 | func NewStringSet(initElements ...string) StringSet {
28 | set := make(map[string]struct{}, len(initElements))
29 | for _, key := range initElements {
30 | set[key] = struct{}{}
31 | }
32 | return set
33 | }
34 |
35 | // Add adds a new key to set
36 | func (s StringSet) Add(key string) {
37 | s[key] = struct{}{}
38 | }
39 |
40 | // AddAll adds a new key to set
41 | func (s StringSet) AddAll(keys ...string) {
42 | for _, key := range keys {
43 | s[key] = struct{}{}
44 | }
45 | }
46 |
47 | // Remove deletes a key from set
48 | func (s StringSet) Remove(key string) {
49 | delete(s, key)
50 | }
51 |
52 | // Contains checks if the set contains the key
53 | func (s StringSet) Contains(key string) bool {
54 | _, ok := s[key]
55 | return ok
56 | }
57 |
58 | // Size returns the size of the set
59 | func (s StringSet) Size() int {
60 | return len(s)
61 | }
62 |
63 | // IsEmpty returns if the size is 0
64 | func (s StringSet) IsEmpty() bool {
65 | return len(s) == 0
66 | }
67 |
--------------------------------------------------------------------------------
/util/set_test.go:
--------------------------------------------------------------------------------
1 | // Copyright (c) 2017 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 util
22 |
23 | import (
24 | "testing"
25 |
26 | "github.com/stretchr/testify/assert"
27 | )
28 |
29 | func TestSet(t *testing.T) {
30 | key := "key"
31 | set := NewStringSet()
32 | assert.Equal(t, 0, set.Size())
33 | assert.True(t, set.IsEmpty())
34 | set.Add(key)
35 | assert.Equal(t, 1, set.Size())
36 | assert.True(t, set.Contains(key))
37 | set.Remove(key)
38 | assert.Equal(t, 0, set.Size())
39 | assert.False(t, set.Contains(key))
40 | set.AddAll("a", "b")
41 | assert.Equal(t, 2, set.Size())
42 | assert.False(t, set.IsEmpty())
43 | }
44 |
--------------------------------------------------------------------------------
/zk/client.go:
--------------------------------------------------------------------------------
1 | // Copyright (c) 2017 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 zk
22 |
23 | import (
24 | "path"
25 | "strconv"
26 | "strings"
27 | "sync"
28 | "time"
29 |
30 | "github.com/pkg/errors"
31 | "github.com/samuel/go-zookeeper/zk"
32 | "github.com/uber-go/go-helix/model"
33 | "github.com/uber-go/tally"
34 | "go.uber.org/zap"
35 | )
36 |
37 | const (
38 | // DefaultSessionTimeout is the time for ZK server to expire client session
39 | // ZK uses 20 times tickTime, which is set to 200 ms
40 | // in existing ZK cluster to ensure proper detection
41 | // of Helix node failures
42 | DefaultSessionTimeout = 4 * time.Second
43 | // FlagsZero is the default ZK data node flag
44 | FlagsZero = int32(0)
45 | // FlagsEphemeral is the ephemeral ZK data node flag
46 | FlagsEphemeral = int32(zk.FlagEphemeral)
47 |
48 | _defaultRetryTimeout = 1 * time.Minute
49 | )
50 |
51 | var (
52 | // ACLPermAll is the mode of ZK nodes where all users have permission to access
53 | ACLPermAll = zk.WorldACL(zk.PermAll)
54 | )
55 |
56 | var (
57 | errOpBeforeConnect = errors.New("zookeeper: called connect() before any ops")
58 | )
59 |
60 | // Connection is the thread safe interface for ZK connection
61 | type Connection interface {
62 | AddAuth(scheme string, auth []byte) error
63 | Children(path string) ([]string, *zk.Stat, error)
64 | ChildrenW(path string) ([]string, *zk.Stat, <-chan zk.Event, error)
65 | Get(path string) ([]byte, *zk.Stat, error)
66 | GetW(path string) ([]byte, *zk.Stat, <-chan zk.Event, error)
67 | Exists(path string) (bool, *zk.Stat, error)
68 | ExistsW(path string) (bool, *zk.Stat, <-chan zk.Event, error)
69 | Set(path string, data []byte, version int32) (*zk.Stat, error)
70 | Create(path string, data []byte, flags int32, acl []zk.ACL) (string, error)
71 | Delete(path string, version int32) error
72 | Multi(ops ...interface{}) ([]zk.MultiResponse, error)
73 | SessionID() int64
74 | SetLogger(zk.Logger)
75 | State() zk.State
76 | Close()
77 | }
78 |
79 | // ConnFactory provides interface that creates ZK connections
80 | type ConnFactory interface {
81 | NewConn() (Connection, <-chan zk.Event, error)
82 | }
83 |
84 | // connFactory creates connections to real/embedded ZK
85 | type connFactory struct {
86 | zkServers []string
87 | sessionTimeout time.Duration
88 | }
89 |
90 | // NewConnFactory creates new connFactory
91 | func NewConnFactory(zkServers []string, sessionTimeout time.Duration) ConnFactory {
92 | return &connFactory{
93 | zkServers: zkServers,
94 | sessionTimeout: sessionTimeout,
95 | }
96 | }
97 |
98 | // NewConn creates new ZK connection to real/embedded ZK
99 | func (f *connFactory) NewConn() (Connection, <-chan zk.Event, error) {
100 | return zk.Connect(f.zkServers, f.sessionTimeout)
101 | }
102 |
103 | // Client wraps utils to communicate with ZK
104 | type Client struct {
105 | logger *zap.Logger
106 | scope tally.Scope
107 |
108 | zkSvr string
109 | sessionTimeout time.Duration
110 | retryTimeout time.Duration
111 | connFactory ConnFactory
112 | zkConn Connection
113 | zkConnMu *sync.RWMutex
114 | // coordinates Go routines waiting on ZK connection events
115 | cond *sync.Cond
116 |
117 | zkEventWatchersMu *sync.RWMutex
118 | zkEventWatchers []Watcher
119 | }
120 |
121 | // Watcher mirrors org.apache.zookeeper.Watcher
122 | type Watcher interface {
123 | Process(e zk.Event)
124 | }
125 |
126 | // ClientOption provides options or ZK client
127 | type ClientOption func(*Client)
128 |
129 | // WithZkSvr configures ZK servers for the client
130 | func WithZkSvr(zkSvr string) ClientOption {
131 | return func(c *Client) {
132 | c.zkSvr = zkSvr
133 | }
134 | }
135 |
136 | // WithSessionTimeout configures sessionTimeout
137 | func WithSessionTimeout(t time.Duration) ClientOption {
138 | return func(c *Client) {
139 | c.sessionTimeout = t
140 | }
141 | }
142 |
143 | // WithRetryTimeout configures retryTimeout for ZK operations
144 | func WithRetryTimeout(t time.Duration) ClientOption {
145 | return func(c *Client) {
146 | c.retryTimeout = t
147 | }
148 | }
149 |
150 | // WithConnFactory configures ConnFactory used to make ZK connections
151 | func WithConnFactory(connFactory ConnFactory) ClientOption {
152 | return func(c *Client) {
153 | c.connFactory = connFactory
154 | }
155 | }
156 |
157 | // NewClient returns new ZK client
158 | func NewClient(logger *zap.Logger, scope tally.Scope, options ...ClientOption) *Client {
159 | mu := &sync.Mutex{}
160 | c := &Client{
161 | cond: sync.NewCond(mu),
162 | retryTimeout: _defaultRetryTimeout,
163 | zkConnMu: &sync.RWMutex{},
164 | zkEventWatchersMu: &sync.RWMutex{},
165 | }
166 | for _, option := range options {
167 | option(c)
168 | }
169 | c.logger = logger.With(zap.String("zkSvr", c.zkSvr))
170 | c.scope = scope.SubScope("helix.zk").Tagged(map[string]string{"zkSvr": c.zkSvr})
171 | if c.connFactory == nil {
172 | zkServers := strings.Split(strings.TrimSpace(c.zkSvr), ",")
173 | c.connFactory = NewConnFactory(zkServers, c.sessionTimeout)
174 | }
175 | return c
176 | }
177 |
178 | // AddWatcher adds a Watcher to zk session event
179 | func (c *Client) AddWatcher(w Watcher) {
180 | c.zkEventWatchersMu.Lock()
181 | c.zkEventWatchers = append(c.zkEventWatchers, w)
182 | c.zkEventWatchersMu.Unlock()
183 | }
184 |
185 | // ClearWatchers removes all the watchers the client has
186 | func (c *Client) ClearWatchers() {
187 | c.zkEventWatchersMu.Lock()
188 | c.zkEventWatchers = nil
189 | c.zkEventWatchersMu.Unlock()
190 | }
191 |
192 | // Connect sets up ZK connection
193 | func (c *Client) Connect() error {
194 | zkConn, eventCh, err := c.connFactory.NewConn()
195 | if err != nil {
196 | return err
197 | }
198 | c.zkConnMu.Lock()
199 | if c.zkConn != nil {
200 | c.zkConn.Close()
201 | }
202 | c.zkConn = zkConn
203 | c.zkConnMu.Unlock()
204 | go c.processEvents(eventCh)
205 | connected := c.waitUntilConnected(c.sessionTimeout)
206 | if !connected {
207 | return errors.New("zookeeper: failed to connect")
208 | }
209 | return nil
210 | }
211 |
212 | func (c *Client) processEvents(eventCh <-chan zk.Event) {
213 | for {
214 | select {
215 | case ev, ok := <-eventCh:
216 | if !ok {
217 | c.logger.Warn("zookeeper has quit, stop processing events")
218 | return
219 | }
220 | switch ev.Type {
221 | case zk.EventSession:
222 | c.logger.Info("receive EventSession", zap.Any("state", ev.State))
223 | c.cond.Broadcast()
224 | c.processSessionEvents(ev)
225 | case zk.EventNotWatching:
226 | c.logger.Info("watchers have been invalidated. zk client should be trying to reconnect")
227 | default:
228 | // no-op since node specific events are handled by caller
229 | }
230 | }
231 | }
232 | }
233 |
234 | func (c *Client) processSessionEvents(ev zk.Event) {
235 | c.zkEventWatchersMu.RLock()
236 | watchers := c.zkEventWatchers
237 | c.zkEventWatchersMu.RUnlock()
238 | for _, watcher := range watchers {
239 | watcher.Process(ev)
240 | }
241 | }
242 |
243 | // IsConnected returns if client is connected to Zookeeper
244 | func (c *Client) IsConnected() bool {
245 | conn := c.getConn()
246 | return conn != nil && conn.State() == zk.StateHasSession
247 | }
248 |
249 | // GetSessionID returns current ZK session ID
250 | func (c *Client) GetSessionID() string {
251 | conn := c.getConn()
252 | if conn == nil {
253 | return ""
254 | }
255 | return strconv.FormatInt(conn.SessionID(), 10)
256 | }
257 |
258 | // Disconnect closes ZK connection
259 | func (c *Client) Disconnect() {
260 | c.ClearWatchers()
261 | conn := c.getConn()
262 | if conn != nil {
263 | conn.Close()
264 | }
265 | }
266 |
267 | // CreateEmptyNode creates an empty node for future use
268 | func (c *Client) CreateEmptyNode(path string) error {
269 | return c.Create(path, []byte(""), FlagsZero, ACLPermAll)
270 | }
271 |
272 | // CreateDataWithPath creates a path with a string
273 | func (c *Client) CreateDataWithPath(p string, data []byte) error {
274 | parent := path.Dir(p)
275 | if err := c.ensurePath(parent); err != nil {
276 | return err
277 | }
278 | return c.Create(p, data, FlagsZero, ACLPermAll)
279 | }
280 |
281 | // Exists checks if a key exists in ZK
282 | func (c *Client) Exists(path string) (bool, *zk.Stat, error) {
283 | var res bool
284 | var stat *zk.Stat
285 | err := c.retryUntilConnected(func() error {
286 | r, s, err := c.getConn().Exists(path)
287 | if err != nil {
288 | return err
289 | }
290 | res = r
291 | stat = s
292 | return nil
293 | })
294 | return res, stat, errors.Wrapf(err, "zk client failed to check existence of %s", path)
295 | }
296 |
297 | // ExistsAll returns if all paths exist
298 | func (c *Client) ExistsAll(paths ...string) (bool, error) {
299 | for _, path := range paths {
300 | if exists, _, err := c.Exists(path); err != nil || !exists {
301 | return false, err
302 | }
303 | }
304 | return true, nil
305 | }
306 |
307 | // Get returns data in ZK path
308 | func (c *Client) Get(path string) ([]byte, *zk.Stat, error) {
309 | var data []byte
310 | var stat *zk.Stat
311 | err := c.retryUntilConnected(func() error {
312 | d, s, err := c.getConn().Get(path)
313 | if err != nil {
314 | return err
315 | }
316 | data = d
317 | stat = s
318 | return nil
319 | })
320 | return data, stat, errors.Wrapf(err, "zk client failed to get data at %s", path)
321 | }
322 |
323 | // GetW returns data in ZK path and watches path
324 | func (c *Client) GetW(path string) ([]byte, <-chan zk.Event, error) {
325 | var data []byte
326 | var events <-chan zk.Event
327 | err := c.retryUntilConnected(func() error {
328 | d, _, evts, err := c.getConn().GetW(path)
329 | if err != nil {
330 | return err
331 | }
332 | data = d
333 | events = evts
334 | return nil
335 | })
336 | return data, events, errors.Wrapf(err, "zk client failed to get and watch data at %s", path)
337 | }
338 |
339 | // Set sets data in ZK path
340 | func (c *Client) Set(path string, data []byte, version int32) error {
341 | err := c.retryUntilConnected(func() error {
342 | _, err := c.getConn().Set(path, data, version)
343 | return err
344 | })
345 | return errors.Wrapf(err, "zk client failed to set data at %s", path)
346 | }
347 |
348 | // SetWithDefaultVersion sets data with default version, -1
349 | func (c *Client) SetWithDefaultVersion(path string, data []byte) error {
350 | return c.Set(path, data, -1)
351 | }
352 |
353 | // Create creates ZK path with data
354 | func (c *Client) Create(path string, data []byte, flags int32, acl []zk.ACL) error {
355 | err := c.retryUntilConnected(func() error {
356 | _, err := c.getConn().Create(path, data, flags, acl)
357 | return err
358 | })
359 | return errors.Wrapf(err, "zk client failed to create data at %s", path)
360 | }
361 |
362 | // Children returns children of ZK path
363 | func (c *Client) Children(path string) ([]string, error) {
364 | var children []string
365 | err := c.retryUntilConnected(func() error {
366 | res, _, err := c.getConn().Children(path)
367 | if err != nil {
368 | return err
369 | }
370 | children = res
371 | return nil
372 | })
373 | return children, errors.Wrapf(err, "zk client failed to get children of %s", path)
374 | }
375 |
376 | // ChildrenW gets children and watches path
377 | func (c *Client) ChildrenW(path string) ([]string, <-chan zk.Event, error) {
378 | children := []string{}
379 | eventCh := make(<-chan zk.Event)
380 |
381 | err := c.retryUntilConnected(func() error {
382 | res, _, evts, err := c.getConn().ChildrenW(path)
383 | if err != nil {
384 | return err
385 | }
386 | children = res
387 | eventCh = evts
388 | return nil
389 | })
390 |
391 | return children, eventCh,
392 | errors.Wrapf(err, "zk client failed to get and watch children of %s", path)
393 | }
394 |
395 | // Delete removes ZK path
396 | func (c *Client) Delete(path string) error {
397 | err := c.retryUntilConnected(func() error {
398 | err := c.getConn().Delete(path, -1)
399 | return err
400 | })
401 | return errors.Wrapf(err, "zk client failed to delete node at %s", path)
402 | }
403 |
404 | // UpdateMapField updates a map field for path
405 | // key is the top-level key in the MapFields
406 | // mapProperty is the inner key
407 | //
408 | // Example:
409 | //
410 | // mapFields":{
411 | // "partition_1":{
412 | // "CURRENT_STATE":"OFFLINE",
413 | // "INFO":""
414 | // }
415 | //
416 | // To set the CURRENT_STATE to ONLINE, use
417 | // UpdateMapField(
418 | // "/CLUSTER/INSTANCES/{instance}/CURRENT_STATE/{sessionID}/{db}",
419 | // "partition_1", "CURRENT_STATE", "ONLINE")
420 | func (c *Client) UpdateMapField(path string, key string, property string, value string) error {
421 | data, stat, err := c.Get(path)
422 | if err != nil {
423 | return err
424 | }
425 |
426 | // convert the result into Message
427 | node, err := model.NewRecordFromBytes(data)
428 | if err != nil {
429 | return err
430 | }
431 |
432 | // update the value
433 | node.SetMapField(key, property, value)
434 |
435 | // marshall to bytes
436 | data, err = node.Marshal()
437 | if err != nil {
438 | return err
439 | }
440 |
441 | // copy back to zookeeper
442 | err = c.Set(path, data, stat.Version)
443 | return err
444 | }
445 |
446 | // UpdateSimpleField updates a simple field
447 | func (c *Client) UpdateSimpleField(path string, key string, value string) error {
448 | data, stat, err := c.Get(path)
449 | if err != nil {
450 | return err
451 | }
452 |
453 | // convert the result into Message
454 | node, err := model.NewRecordFromBytes(data)
455 | if err != nil {
456 | return err
457 | }
458 |
459 | // update the value
460 | node.SetSimpleField(key, value)
461 |
462 | // marshall to bytes
463 | data, err = node.Marshal()
464 | if err != nil {
465 | return err
466 | }
467 |
468 | // copy back to zookeeper
469 | err = c.Set(path, data, stat.Version)
470 | return err
471 | }
472 |
473 | // GetSimpleFieldValueByKey returns value in simple field by key
474 | func (c *Client) GetSimpleFieldValueByKey(path string, key string) (string, error) {
475 | data, stat, err := c.Get(path)
476 | if err != nil {
477 | return "", err
478 | }
479 |
480 | record, err := model.NewRecordFromBytes(data)
481 | if err != nil {
482 | return "", err
483 | }
484 | record.Version = stat.Version
485 |
486 | v, ok := record.GetSimpleField(key)
487 | if !ok {
488 | return "", nil
489 | }
490 | return v, nil
491 | }
492 |
493 | // DeleteTree removes ZK path and its children
494 | func (c *Client) DeleteTree(path string) error {
495 | if exists, _, err := c.Exists(path); !exists || err != nil {
496 | return err
497 | }
498 |
499 | children, err := c.Children(path)
500 | if err != nil {
501 | return err
502 | }
503 |
504 | if len(children) == 0 {
505 | err := c.Delete(path)
506 | return err
507 | }
508 |
509 | for _, child := range children {
510 | p := path + "/" + child
511 | e := c.DeleteTree(p)
512 | if e != nil {
513 | return e
514 | }
515 | }
516 |
517 | return c.Delete(path)
518 | }
519 |
520 | // RemoveMapFieldKey removes a map field by key
521 | func (c *Client) RemoveMapFieldKey(path string, key string) error {
522 | data, stat, err := c.Get(path)
523 | if err != nil {
524 | return err
525 | }
526 |
527 | node, err := model.NewRecordFromBytes(data)
528 | if err != nil {
529 | return err
530 | }
531 |
532 | node.RemoveMapField(key)
533 |
534 | data, err = node.Marshal()
535 | if err != nil {
536 | return err
537 | }
538 |
539 | // save the data back to zookeeper
540 | err = c.Set(path, data, stat.Version)
541 | return err
542 | }
543 |
544 | // GetRecordFromPath returns message by ZK path
545 | func (c *Client) GetRecordFromPath(path string) (*model.ZNRecord, error) {
546 | data, stat, err := c.Get(path)
547 | if err != nil {
548 | return nil, err
549 | }
550 | record, err := model.NewRecordFromBytes(data)
551 | if err != nil {
552 | return nil, err
553 | }
554 | record.Version = stat.Version
555 | return record, nil
556 | }
557 |
558 | // SetDataForPath updates data at given ZK path
559 | func (c *Client) SetDataForPath(path string, data []byte, version int32) error {
560 | return c.Set(path, data, version)
561 | }
562 |
563 | // SetRecordForPath sets a record in give ZK path
564 | func (c *Client) SetRecordForPath(path string, r *model.ZNRecord) error {
565 | version, err := c.getVersionFromPath(path)
566 | if err != nil {
567 | return err
568 | }
569 |
570 | data, err := r.Marshal()
571 | if err != nil {
572 | return err
573 | }
574 | return c.SetDataForPath(path, data, version)
575 | }
576 |
577 | func (c *Client) getConn() Connection {
578 | c.zkConnMu.RLock()
579 | defer c.zkConnMu.RUnlock()
580 | return c.zkConn
581 | }
582 |
583 | func (c *Client) getVersionFromPath(p string) (int32, error) {
584 | exists, s, _ := c.Exists(p)
585 | if !exists {
586 | version := int32(0)
587 | err := c.ensurePath(p)
588 | if err != nil {
589 | return version, errors.Wrap(err, "failed to get version from path: %v")
590 | }
591 | return version, nil
592 | }
593 | return s.Version, nil
594 | }
595 |
596 | // EnsurePath makes sure the specified path exists.
597 | // If not, create it
598 | func (c *Client) ensurePath(p string) error {
599 | exists, _, err := c.Exists(p)
600 | if err != nil {
601 | return err
602 | }
603 | if exists {
604 | return nil
605 | }
606 | err = c.ensurePath(path.Dir(p))
607 | if err != nil {
608 | return err
609 | }
610 | return c.CreateEmptyNode(p)
611 | }
612 |
613 | // Mirrors org.I0Itec.zkclient.Client#retryUntilConnected
614 | func (c *Client) retryUntilConnected(fn func() error) error {
615 | conn := c.getConn()
616 | if conn == nil {
617 | return errOpBeforeConnect
618 | }
619 | startTime := time.Now()
620 | for {
621 | if conn.State() == zk.StateDisconnected {
622 | return errors.New("zookeeper: client disconnected")
623 | }
624 | if (time.Since(startTime)) > c.retryTimeout {
625 | return errors.New("zookeeper: retry has timed out")
626 | }
627 | err := fn()
628 | if err == zk.ErrConnectionClosed || err == zk.ErrSessionExpired {
629 | c.waitUntilConnected(c.retryTimeout)
630 | continue
631 | } else if err != nil {
632 | return errors.Wrap(err, "zookeeper error shouldn't be retried")
633 | }
634 | return nil
635 | }
636 | }
637 |
638 | // waitUntilConnected returns if the client is connected in one of two scenarios
639 | // 1) client is connected within duration t
640 | // 2) client is not connected and duration t is exhausted
641 | // ZooKeeper state change is broadcast by the cond variable of the client,
642 | // note c.cond.L is not locked when Wait first resumes
643 | func (c *Client) waitUntilConnected(t time.Duration) bool {
644 | startTime := time.Now()
645 | for !c.IsConnected() {
646 | if time.Since(startTime) > t {
647 | return false
648 | }
649 | doneCh := make(chan struct{})
650 | c.cond.L.Lock()
651 | go func() {
652 | c.cond.Wait()
653 | c.cond.L.Unlock()
654 | close(doneCh)
655 | }()
656 | select {
657 | case <-time.After(t):
658 | case <-doneCh:
659 | }
660 | }
661 | return true
662 | }
663 |
--------------------------------------------------------------------------------
/zk/client_test.go:
--------------------------------------------------------------------------------
1 | // Copyright (c) 2017 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 zk
22 |
23 | import (
24 | "fmt"
25 | "math/rand"
26 | "path"
27 | "sync"
28 | "testing"
29 | "time"
30 |
31 | "github.com/pkg/errors"
32 | "github.com/samuel/go-zookeeper/zk"
33 | "github.com/stretchr/testify/suite"
34 | "github.com/uber-go/go-helix/model"
35 | "github.com/uber-go/tally"
36 | "go.uber.org/zap"
37 | )
38 |
39 | type ZKClientTestSuite struct {
40 | BaseZkTestSuite
41 |
42 | zkClient *Client
43 | }
44 |
45 | type CountEventWatcher struct {
46 | count int
47 | sync.RWMutex
48 | }
49 |
50 | func (w *CountEventWatcher) Process(e zk.Event) {
51 | w.Lock()
52 | defer w.Unlock()
53 | w.count++
54 | }
55 |
56 | func (w *CountEventWatcher) GetCount() int {
57 | w.RLock()
58 | defer w.RUnlock()
59 | return w.count
60 | }
61 |
62 | func TestZKClientTestSuite(t *testing.T) {
63 | s := &ZKClientTestSuite{}
64 | suite.Run(t, s)
65 | }
66 |
67 | func (s *ZKClientTestSuite) SetupTest() {
68 | s.zkClient = NewClient(zap.NewNop(), tally.NoopScope, WithZkSvr(s.ZkConnectString),
69 | WithSessionTimeout(DefaultSessionTimeout))
70 | }
71 |
72 | func (s *ZKClientTestSuite) TearDownTest() {
73 | if s.zkClient.IsConnected() {
74 | s.zkClient.Disconnect()
75 | }
76 | }
77 |
78 | func (s *ZKClientTestSuite) TestEmbeddedZk() {
79 | err := EnsureZookeeperUp(s.EmbeddedZkPath)
80 | s.NoError(err)
81 | err = StopZookeeper(s.EmbeddedZkPath)
82 | s.NoError(err)
83 | err = EnsureZookeeperUp(s.EmbeddedZkPath)
84 | s.NoError(err)
85 | }
86 |
87 | func (s *ZKClientTestSuite) TestZKConnectAndDisconnect() {
88 | s.False(s.zkClient.IsConnected())
89 |
90 | err := s.zkClient.Connect()
91 | s.NoError(err)
92 | s.True(s.zkClient.IsConnected())
93 | s.True(len(s.zkClient.GetSessionID()) > 0)
94 |
95 | s.zkClient.Disconnect()
96 | // time.Sleep is needed because ZK client disconnects asynchronously,
97 | // which is in parity with the Java client
98 | time.Sleep(time.Millisecond)
99 | s.False(s.zkClient.IsConnected())
100 | }
101 |
102 | func (s *ZKClientTestSuite) TestBasicZkOps() {
103 | testData := fmt.Sprintf("%d", rand.Int())
104 | testData1 := fmt.Sprintf("%d", rand.Int())
105 | testPath := fmt.Sprintf("/%d/%d", rand.Int63(), rand.Int63())
106 | testPath1 := fmt.Sprintf("/%d/%d", rand.Int63(), rand.Int63())
107 |
108 | err := s.zkClient.CreateDataWithPath(testPath, []byte(testData))
109 | s.Equal(errOpBeforeConnect, errors.Cause(err))
110 | err = s.zkClient.Connect()
111 | s.NoError(err)
112 | s.True(s.zkClient.IsConnected())
113 | err = s.zkClient.CreateDataWithPath(testPath, []byte(testData))
114 | s.NoError(err)
115 | bytes, stat, err := s.zkClient.Get(testPath)
116 | s.NoError(err)
117 | s.Equal(testData, string(bytes))
118 | s.Equal(int32(0), stat.Version)
119 | err = s.zkClient.SetWithDefaultVersion(testPath, []byte(testData1))
120 | s.NoError(err)
121 | bytes, eventCh, err := s.zkClient.GetW(testPath)
122 | s.NoError(err)
123 | s.Equal(testData1, string(bytes))
124 | err = s.zkClient.Set(testPath, []byte(testData), 1)
125 | ev := <-eventCh
126 | s.Equal(zk.EventNodeDataChanged, ev.Type)
127 |
128 | parent := path.Dir(testPath)
129 | paths, eventCh, err := s.zkClient.ChildrenW(parent)
130 | s.NoError(err)
131 | s.Equal(1, len(paths))
132 | s.zkClient.Delete(testPath)
133 | ev = <-eventCh
134 | s.Equal(zk.EventNodeChildrenChanged, ev.Type)
135 | exists, _, err := s.zkClient.Exists(parent)
136 | s.True(exists)
137 | err = s.zkClient.CreateDataWithPath(testPath, []byte(testData))
138 | s.NoError(err)
139 | err = s.zkClient.DeleteTree(parent)
140 | s.NoError(err)
141 |
142 | err = s.zkClient.CreateDataWithPath(testPath, []byte(testData))
143 | s.NoError(err)
144 | s.False(s.zkClient.ExistsAll(testPath, testPath1))
145 | err = s.zkClient.CreateDataWithPath(testPath1, []byte(testData))
146 | s.NoError(err)
147 | s.True(s.zkClient.ExistsAll(testPath, testPath1))
148 | }
149 |
150 | func (s *ZKClientTestSuite) TestHelixRecordOps() {
151 | path := fmt.Sprintf("/%d", rand.Int63())
152 | partition := fmt.Sprintf("partition_%d", rand.Int())
153 | state := "ONLINE"
154 | key := fmt.Sprintf("%d", rand.Int())
155 | value := fmt.Sprintf("%d", rand.Int())
156 | c := NewClient(zap.NewNop(), tally.NoopScope, WithZkSvr(s.ZkConnectString),
157 | WithSessionTimeout(DefaultSessionTimeout))
158 | err := c.Connect()
159 | s.NoError(err)
160 |
161 | record := &model.ZNRecord{}
162 | err = c.SetRecordForPath(path, record)
163 | s.NoError(err)
164 | err = c.UpdateMapField(path, partition, model.FieldKeyCurrentState, state)
165 | s.NoError(err)
166 | err = c.UpdateSimpleField(path, key, value)
167 | record, err = c.GetRecordFromPath(path)
168 | s.NoError(err)
169 | s.Equal(state, record.GetMapField(partition, model.FieldKeyCurrentState))
170 | val, err := c.GetSimpleFieldValueByKey(path, key)
171 | s.NoError(err)
172 | s.Equal(value, val)
173 | err = c.RemoveMapFieldKey(path, partition)
174 | s.NoError(err)
175 | record, err = c.GetRecordFromPath(path)
176 | s.NoError(err)
177 | s.Equal("", record.GetMapField(partition, model.FieldKeyCurrentState))
178 | c.Disconnect()
179 | }
180 |
181 | func (s *ZKClientTestSuite) TestWatcher() {
182 | c := NewClient(zap.NewNop(), tally.NoopScope, WithZkSvr(s.ZkConnectString),
183 | WithSessionTimeout(DefaultSessionTimeout))
184 | watcher := &CountEventWatcher{}
185 | c.AddWatcher(watcher)
186 | err := c.Connect()
187 | s.NoError(err)
188 | s.Equal(3, watcher.GetCount())
189 | c.Disconnect()
190 | s.Equal(3, watcher.GetCount())
191 | }
192 |
193 | // TestRetryUntilConnected func failed once then succeeds on retry if the client is connected
194 | func (s *ZKClientTestSuite) TestRetryUntilConnected() {
195 | z := NewFakeZk()
196 | client := s.createClientWithFakeConn(z)
197 | invokeCounter := 0
198 | client.Connect()
199 | s.Equal(1, len(z.GetConnections()))
200 | z.SetState(client.zkConn, zk.StateHasSession)
201 | s.NoError(client.retryUntilConnected(getFailOnceFunc(&invokeCounter)))
202 | s.Equal(2, invokeCounter)
203 | z.stop()
204 | }
205 |
206 | // TestRetryUntilConnectedWithoutSignal func fails and times out if client is not connected and
207 | // condition receives no signals
208 | func (s *ZKClientTestSuite) TestRetryUntilConnectedWithoutSignal() {
209 | z := NewFakeZk()
210 | client := s.createClientWithFakeConn(z)
211 | invokeCounter := 0
212 | expiringFn := func() error {
213 | invokeCounter++
214 | return zk.ErrSessionExpired
215 | }
216 | client.Connect()
217 | z.SetState(client.zkConn, zk.StateConnecting)
218 | s.Error(client.retryUntilConnected(expiringFn))
219 | s.Equal(1, invokeCounter)
220 |
221 | }
222 |
223 | // TestRetryUntilConnectedWithSignal func fails once, then succeeds on retry if the condition
224 | // receives a signal
225 | func (s *ZKClientTestSuite) TestRetryUntilConnectedWithSignal() {
226 | z := NewFakeZk()
227 | client := s.createClientWithFakeConn(z)
228 | client.Connect()
229 | z.SetState(client.zkConn, zk.StateConnecting)
230 | go func() {
231 | z.SetState(client.zkConn, zk.StateHasSession)
232 | // wait to allow event loop process op
233 | time.Sleep(10 * time.Millisecond)
234 | client.cond.Broadcast()
235 | }()
236 | invokeCounter := 0
237 | s.NoError(client.retryUntilConnected(getFailOnceFunc(&invokeCounter)))
238 | s.Equal(2, invokeCounter)
239 | z.stop()
240 | }
241 |
242 | // TestZkSizeLimit ensures zk write/update requests fail if size exceeds 1MB
243 | func (s *ZKClientTestSuite) TestZkSizeLimit() {
244 | client := s.CreateAndConnectClient()
245 | defer client.Disconnect()
246 | numMb := 1024 * 1024
247 | legalData := make([]byte, numMb-128)
248 | legalData2 := make([]byte, numMb-128)
249 | legalData2[0] = 1
250 | oversizeData := make([]byte, numMb)
251 | path := s.createRandomPath()
252 |
253 | // create fails if data exceeds 1MB
254 | err := client.Create(path, oversizeData, FlagsZero, ACLPermAll)
255 | s.Error(err)
256 |
257 | // create succeeds if data size is legal
258 | err = client.Create(path, legalData, FlagsZero, ACLPermAll)
259 | s.Nil(err)
260 | res, _, err := client.Get(path)
261 | s.Equal(res, legalData)
262 |
263 | // set fails if data exceeds 1MB
264 | err = client.Set(path, oversizeData, -1)
265 | s.Error(err)
266 |
267 | // set succeeds if data size is legal
268 | err = client.Set(path, legalData2, -1)
269 | s.Nil(err)
270 | res, _, err = client.Get(path)
271 | s.Equal(res, legalData2)
272 | }
273 |
274 | func (s *ZKClientTestSuite) createClientWithFakeConn(z *FakeZk) *Client {
275 | return NewClient(zap.NewNop(), tally.NoopScope, WithConnFactory(z), WithRetryTimeout(time.Second))
276 | }
277 |
278 | func (s ZKClientTestSuite) createRandomPath() string {
279 | return fmt.Sprintf("/%d", rand.Int63())
280 | }
281 |
282 | // getFailOnceFunc returns a function that succeeds on second invocation
283 | func getFailOnceFunc(invokeCounter *int) func() error {
284 | shouldPass := false
285 | return func() error {
286 | *invokeCounter++
287 | if shouldPass {
288 | return nil
289 | }
290 | shouldPass = true
291 | return zk.ErrSessionExpired
292 | }
293 | }
294 |
--------------------------------------------------------------------------------
/zk/embedded/start.sh:
--------------------------------------------------------------------------------
1 | #!/bin/bash
2 |
3 | pidc=`jps -v | grep embedded-zookeeper | wc -l`
4 | if [[ $pidc -ne 0 ]] ; then
5 | echo "test zookeeper already running"
6 | else
7 | echo "starting test zookeeper"
8 | mkdir -p zookeeper-data
9 | java -Dname=embedded-zookeeper -jar zookeeper-3.4.9-fatjar.jar server zk.cfg > zookeeper-data/zookeeper.log &
10 | fi
11 |
--------------------------------------------------------------------------------
/zk/embedded/stop.sh:
--------------------------------------------------------------------------------
1 | #!/bin/bash
2 |
3 | echo "killing test zookeeper"
4 | jps -v | grep embedded-zookeeper | awk '{print $1}' | xargs kill -9
5 | rm -rf zookeeper-data
6 |
--------------------------------------------------------------------------------
/zk/embedded/zk.cfg:
--------------------------------------------------------------------------------
1 | tickTime=2000
2 | dataDir=zookeeper-data
3 | clientPort=2181
4 |
--------------------------------------------------------------------------------
/zk/embedded/zookeeper-3.4.9-fatjar.jar:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/uber-go/go-helix/d9281c27d5db32cdd0f14e79e4782f23b505a7d0/zk/embedded/zookeeper-3.4.9-fatjar.jar
--------------------------------------------------------------------------------
/zk/fake_zk.go:
--------------------------------------------------------------------------------
1 | // Copyright (c) 2017 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 zk
22 |
23 | import (
24 | "fmt"
25 | "github.com/samuel/go-zookeeper/zk"
26 | )
27 |
28 | type connInfo struct {
29 | state zk.State
30 | eventCh chan zk.Event
31 | }
32 |
33 | // FakeZk provides utility to make fake connections and manipulate connection states
34 | type FakeZk struct {
35 | connToConnInfo map[Connection]*connInfo
36 | opChan chan interface{}
37 |
38 | defaultConnectionState zk.State
39 | }
40 |
41 | // FakeZkOption is the optional arg to create a FakeZk
42 | type FakeZkOption func(*FakeZk)
43 |
44 | // DefaultConnectionState sets the default state when a fake connection is made
45 | func DefaultConnectionState(state zk.State) FakeZkOption {
46 | return func(zk *FakeZk) {
47 | zk.defaultConnectionState = state
48 | }
49 | }
50 |
51 | // NewFakeZk creates new FakeZk test utility
52 | func NewFakeZk(opts ...FakeZkOption) *FakeZk {
53 | z := &FakeZk{
54 | connToConnInfo: map[Connection]*connInfo{},
55 | opChan: make(chan interface{}, 1),
56 | defaultConnectionState: zk.StateDisconnected,
57 | }
58 | for _, opt := range opts {
59 | opt(z)
60 | }
61 | go z.run()
62 | return z
63 | }
64 |
65 | // NewConn makes new fake ZK connection
66 | func (z *FakeZk) NewConn() (Connection, <-chan zk.Event, error) {
67 | respCh := make(chan makeConnResp, 1)
68 | z.opChan <- makeConnReq{
69 | c: respCh,
70 | }
71 | resp := <-respCh
72 | return resp.conn, resp.eventCh, resp.err
73 | }
74 |
75 | // GetState returns state by ZK connection
76 | func (z *FakeZk) GetState(conn Connection) zk.State {
77 | respCh := make(chan zk.State, 1)
78 | z.opChan <- getConnStateReq{
79 | c: respCh,
80 | conn: conn,
81 | }
82 | return <-respCh
83 | }
84 |
85 | // SetState sets state of ZK connection
86 | func (z *FakeZk) SetState(conn Connection, state zk.State) {
87 | respCh := make(chan struct{}, 1)
88 | z.opChan <- setConnStateReq{
89 | eType: zk.EventSession,
90 | state: state,
91 | c: respCh,
92 | conn: conn,
93 | }
94 | <-respCh
95 | }
96 |
97 | // GetConnections returns all of the connections FakeZk has made
98 | func (z *FakeZk) GetConnections() []*FakeZkConn {
99 | var result []*FakeZkConn
100 | for connection := range z.connToConnInfo {
101 | result = append(result, connection.(*FakeZkConn))
102 | }
103 | return result
104 | }
105 |
106 | func (z *FakeZk) run() {
107 | for op := range z.opChan {
108 | z.performOp(op)
109 | }
110 | }
111 |
112 | func (z *FakeZk) stop() {
113 | close(z.opChan)
114 | }
115 |
116 | func (z *FakeZk) performOp(op interface{}) {
117 | switch op := op.(type) {
118 | case makeConnReq:
119 | z.makeConn(op)
120 | case setConnStateReq:
121 | z.setConnState(op)
122 | case getConnStateReq:
123 | z.getConnState(op)
124 | default:
125 | panic(fmt.Sprintf("fake zk received unknown op %v", op))
126 | }
127 | }
128 |
129 | func (z *FakeZk) makeConn(op makeConnReq) {
130 | eventCh := make(chan zk.Event)
131 | conn := NewFakeZkConn(z)
132 | z.connToConnInfo[conn] = &connInfo{
133 | state: z.defaultConnectionState,
134 | eventCh: eventCh,
135 | }
136 | resp := makeConnResp{
137 | conn: conn,
138 | eventCh: eventCh,
139 | err: nil,
140 | }
141 | op.c <- resp
142 | }
143 |
144 | func (z *FakeZk) getConnState(op getConnStateReq) {
145 | connInfo, ok := z.connToConnInfo[op.conn]
146 | if !ok {
147 | panic(fmt.Sprintf("fake zk has no connection for op %+v", op))
148 | }
149 | op.c <- connInfo.state
150 | }
151 |
152 | func (z *FakeZk) setConnState(op setConnStateReq) {
153 | connInfo, ok := z.connToConnInfo[op.conn]
154 | if !ok {
155 | panic(fmt.Sprintf("fake zk has no connection for op %+v", op))
156 | }
157 | connInfo.state = op.state
158 | connInfo.eventCh <- zk.Event{State: op.state, Type: op.eType}
159 | if op.state == zk.StateExpired {
160 | op.conn.(*FakeZkConn).invalidateWatchers(zk.ErrSessionExpired)
161 | }
162 | op.c <- struct{}{}
163 | }
164 |
165 | type makeConnReq struct {
166 | c chan makeConnResp
167 | }
168 |
169 | type makeConnResp struct {
170 | conn Connection
171 | eventCh <-chan zk.Event
172 | err error
173 | }
174 |
175 | type getConnStateReq struct {
176 | conn Connection
177 | c chan zk.State
178 | }
179 |
180 | type setConnStateReq struct {
181 | eType zk.EventType
182 | state zk.State
183 | conn Connection
184 | c chan struct{}
185 | }
186 |
--------------------------------------------------------------------------------
/zk/fake_zk_conn.go:
--------------------------------------------------------------------------------
1 | // Copyright (c) 2017 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 zk
22 |
23 | import (
24 | "sync"
25 |
26 | "github.com/samuel/go-zookeeper/zk"
27 | )
28 |
29 | // FakeZkConn is a fake ZK connection for testing
30 | type FakeZkConn struct {
31 | pathWatchers
32 | history *MethodCallHistory
33 | zk *FakeZk
34 | }
35 |
36 | // NewFakeZkConn creates a FakeZkConn
37 | func NewFakeZkConn(zk *FakeZk) *FakeZkConn {
38 | return &FakeZkConn{
39 | zk: zk,
40 | history: &MethodCallHistory{
41 | dict: make(map[string][]*MethodCall),
42 | },
43 | }
44 | }
45 |
46 | // AddAuth addds auth info
47 | func (c *FakeZkConn) AddAuth(scheme string, auth []byte) error {
48 | c.history.addToHistory("AddAuth", scheme, auth)
49 | return nil
50 | }
51 |
52 | // Children returns children of a path
53 | func (c *FakeZkConn) Children(path string) ([]string, *zk.Stat, error) {
54 | c.history.addToHistory("Children", path)
55 | return nil, nil, nil
56 | }
57 |
58 | // ChildrenW returns children and watcher channel of a path
59 | func (c *FakeZkConn) ChildrenW(path string) ([]string, *zk.Stat, <-chan zk.Event, error) {
60 | eventCh := c.addWatcher()
61 | c.history.addToHistory("ChildrenW", path)
62 | return nil, nil, eventCh, nil
63 | }
64 |
65 | // Get returns node by path
66 | func (c *FakeZkConn) Get(path string) ([]byte, *zk.Stat, error) {
67 | c.history.addToHistory("Get", path)
68 | return nil, nil, nil
69 | }
70 |
71 | // GetW returns node and watcher channel of path
72 | func (c *FakeZkConn) GetW(path string) ([]byte, *zk.Stat, <-chan zk.Event, error) {
73 | eventCh := c.addWatcher()
74 | c.history.addToHistory("GetW", path)
75 | return nil, nil, eventCh, nil
76 | }
77 |
78 | // Exists returns if the path exists
79 | func (c *FakeZkConn) Exists(path string) (bool, *zk.Stat, error) {
80 | c.history.addToHistory("Exists", path)
81 | return true, nil, nil
82 | }
83 |
84 | // ExistsW returns if path exists and watcher chan of path
85 | func (c *FakeZkConn) ExistsW(path string) (bool, *zk.Stat, <-chan zk.Event, error) {
86 | eventCh := c.addWatcher()
87 | c.history.addToHistory("ExistsW", path)
88 | return false, nil, eventCh, nil
89 | }
90 |
91 | // Set sets data for path
92 | func (c *FakeZkConn) Set(path string, data []byte, version int32) (*zk.Stat, error) {
93 | c.history.addToHistory("Set", path, data, version)
94 | return nil, nil
95 | }
96 |
97 | // Create creates new ZK node
98 | func (c *FakeZkConn) Create(path string, data []byte, flags int32, acl []zk.ACL) (string, error) {
99 | c.history.addToHistory("Create", path, data, flags, acl)
100 | return "", nil
101 | }
102 |
103 | // Delete deletes ZK node
104 | func (c *FakeZkConn) Delete(path string, version int32) error {
105 | c.history.addToHistory("Delete", path, version)
106 | return nil
107 | }
108 |
109 | // Multi executes multiple ZK operations
110 | func (c *FakeZkConn) Multi(ops ...interface{}) ([]zk.MultiResponse, error) {
111 | c.history.addToHistory("Multi", ops)
112 | return nil, nil
113 | }
114 |
115 | // SessionID returns session ID
116 | func (c *FakeZkConn) SessionID() int64 {
117 | c.history.addToHistory("SessionID")
118 | return int64(0)
119 | }
120 |
121 | // SetLogger sets loggeer for the client
122 | func (c *FakeZkConn) SetLogger(l zk.Logger) {
123 | c.history.addToHistory("SetLogger", l)
124 | }
125 |
126 | // State returns state of the client
127 | func (c *FakeZkConn) State() zk.State {
128 | c.history.addToHistory("State")
129 | return c.zk.GetState(c)
130 | }
131 |
132 | // Close closes the connection to ZK
133 | func (c *FakeZkConn) Close() {
134 | c.invalidateWatchers(zk.ErrClosing)
135 | c.history.addToHistory("Close")
136 | }
137 |
138 | // GetHistory returns history
139 | func (c *FakeZkConn) GetHistory() *MethodCallHistory {
140 | return c.history
141 | }
142 |
143 | type pathWatchers struct {
144 | sync.RWMutex
145 | watchers []chan zk.Event
146 | }
147 |
148 | func (p *pathWatchers) addWatcher() chan zk.Event {
149 | eventCh := make(chan zk.Event)
150 | p.Lock()
151 | p.watchers = append(p.watchers, eventCh)
152 | p.Unlock()
153 | return eventCh
154 | }
155 |
156 | func (p *pathWatchers) invalidateWatchers(err error) {
157 | ev := zk.Event{Type: zk.EventNotWatching, State: zk.StateDisconnected, Err: err}
158 | p.Lock()
159 | for _, watcher := range p.watchers {
160 | watcher <- ev
161 | close(watcher)
162 | }
163 | p.watchers = nil
164 | p.Unlock()
165 | }
166 |
167 | // MethodCall represents a call record
168 | type MethodCall struct {
169 | MethodName string
170 | Params []interface{}
171 | }
172 |
173 | // MethodCallHistory represents the history of the method called on the connection
174 | type MethodCallHistory struct {
175 | sync.RWMutex
176 | history []*MethodCall
177 | dict map[string][]*MethodCall
178 | }
179 |
180 | func (h *MethodCallHistory) addToHistory(method string, params ...interface{}) {
181 | methodCall := &MethodCall{method, params}
182 | h.Lock()
183 | h.history = append(h.history, methodCall)
184 | h.dict[method] = append(h.dict[method], methodCall)
185 | h.Unlock()
186 | }
187 |
188 | // GetHistory returns all of the histories
189 | func (h *MethodCallHistory) GetHistory() []*MethodCall {
190 | defer h.RUnlock()
191 | h.RLock()
192 | return h.history
193 | }
194 |
195 | // GetHistoryForMethod returns all of the histories of a method
196 | func (h *MethodCallHistory) GetHistoryForMethod(method string) []*MethodCall {
197 | defer h.RUnlock()
198 | h.RLock()
199 | return h.dict[method]
200 | }
201 |
--------------------------------------------------------------------------------
/zk/fake_zk_conn_test.go:
--------------------------------------------------------------------------------
1 | // Copyright (c) 2017 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 zk
22 |
23 | import (
24 | "testing"
25 |
26 | "github.com/stretchr/testify/assert"
27 | )
28 |
29 | var (
30 | _defaultPath = "/"
31 | _defaultVersion = int32(0)
32 | )
33 |
34 | func TestFakeZkConn(t *testing.T) {
35 | conn := NewFakeZkConn(NewFakeZk())
36 | err := conn.AddAuth(_defaultPath, nil)
37 | assert.NoError(t, err)
38 | _, _, err = conn.Children(_defaultPath)
39 | assert.NoError(t, err)
40 | _, _, childrenCh, err := conn.ChildrenW(_defaultPath)
41 | assert.NoError(t, err)
42 | _, _, err = conn.Get(_defaultPath)
43 | assert.NoError(t, err)
44 | _, _, getCh, err := conn.GetW(_defaultPath)
45 | assert.NoError(t, err)
46 | _, _, err = conn.Exists(_defaultPath)
47 | assert.NoError(t, err)
48 | _, _, existsCh, err := conn.ExistsW(_defaultPath)
49 | assert.NoError(t, err)
50 | _, err = conn.Set(_defaultPath, nil, _defaultVersion)
51 | assert.NoError(t, err)
52 | conn.Create(_defaultPath, nil, _defaultVersion, nil)
53 | assert.NoError(t, err)
54 | err = conn.Delete(_defaultPath, _defaultVersion)
55 | assert.NoError(t, err)
56 | _, err = conn.Multi()
57 | assert.NoError(t, err)
58 | conn.SessionID()
59 | conn.SetLogger(nil)
60 |
61 | // verify history
62 | calls := conn.GetHistory().GetHistoryForMethod("AddAuth")
63 | assert.Equal(t, 1, len(calls))
64 | allHistory := conn.GetHistory().GetHistory()
65 | assert.Equal(t, 13, len(allHistory))
66 |
67 | // close and cleanup
68 | doneCh := make(chan struct{})
69 | go func() {
70 | for {
71 | select {
72 | case <-childrenCh:
73 | case <-getCh:
74 | case <-existsCh:
75 | case <-doneCh:
76 | break
77 | }
78 | }
79 | }()
80 | conn.Close()
81 | doneCh <- struct{}{}
82 | }
83 |
--------------------------------------------------------------------------------
/zk/test_util.go:
--------------------------------------------------------------------------------
1 | // Copyright (c) 2017 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 zk
22 |
23 | import (
24 | "bufio"
25 | "errors"
26 | "log"
27 | "os"
28 | "os/exec"
29 | "path"
30 | "sync"
31 | "time"
32 |
33 | "github.com/samuel/go-zookeeper/zk"
34 | "github.com/stretchr/testify/suite"
35 | "github.com/uber-go/tally"
36 | "go.uber.org/zap"
37 | )
38 |
39 | // EmbeddedZkServer is the connect string for embedded ZK server
40 | const EmbeddedZkServer = "localhost:2181"
41 |
42 | var mu = sync.Mutex{}
43 |
44 | // BaseZkTestSuite provides utility to test Zookeeper functions without Helix admin
45 | type BaseZkTestSuite struct {
46 | suite.Suite
47 |
48 | EmbeddedZkPath string
49 | ZkConnectString string
50 | }
51 |
52 | // SetupSuite ensures ZK server is up
53 | func (s *BaseZkTestSuite) SetupSuite() {
54 | s.ZkConnectString = EmbeddedZkServer
55 |
56 | if s.EmbeddedZkPath == "" {
57 | s.EmbeddedZkPath = path.Join(os.Getenv("APP_ROOT"), "zk/embedded")
58 | }
59 | err := EnsureZookeeperUp(s.EmbeddedZkPath)
60 | s.NoError(err)
61 | }
62 |
63 | // CreateAndConnectClient creates ZK client and connects to ZK server
64 | func (s *BaseZkTestSuite) CreateAndConnectClient() *Client {
65 | zkClient := NewClient(
66 | zap.NewNop(), tally.NoopScope, WithZkSvr(s.ZkConnectString), WithSessionTimeout(time.Second))
67 | err := zkClient.Connect()
68 | s.NoError(err)
69 | return zkClient
70 | }
71 |
72 | // EnsureZookeeperUp starts the embedded (test) Zookeeper if not running.
73 | func EnsureZookeeperUp(scriptRelativeDirPath string) error {
74 | mu.Lock()
75 | defer mu.Unlock()
76 |
77 | if isEmbeddedZookeeperStarted(3 * time.Second) {
78 | return nil
79 | }
80 |
81 | err := startEmbeddedZookeeper(scriptRelativeDirPath)
82 | if err != nil {
83 | log.Println("Unable to start Zookeeper server: ", err)
84 | return err
85 | }
86 |
87 | log.Println("Zookeeper server is up.")
88 |
89 | return nil
90 | }
91 |
92 | // StopZookeeper stops the embedded (test) Zookeeper if running.
93 | func StopZookeeper(scriptRelativeDirPath string) error {
94 | mu.Lock()
95 | defer mu.Unlock()
96 |
97 | _, _, err := zk.Connect([]string{EmbeddedZkServer}, time.Second)
98 | if err != nil {
99 | return nil
100 | }
101 |
102 | err = stopEmbeddedZookeeper(scriptRelativeDirPath)
103 | if err != nil {
104 | log.Println("Unable to stop Zookeeper server: ", err)
105 | return err
106 | }
107 | log.Println("Zookeeper server is stopped.")
108 | return nil
109 | }
110 |
111 | func isEmbeddedZookeeperStarted(timeout time.Duration) bool {
112 | zkConn, _, err := zk.Connect([]string{EmbeddedZkServer}, time.Second)
113 | if err == nil && zkConn != nil {
114 | defer zkConn.Close()
115 | done := time.After(timeout)
116 | loop:
117 | for { // zk.Connect is async
118 | if zkConn.State() == zk.StateHasSession {
119 | return true
120 | }
121 | select {
122 | case <-done:
123 | break loop
124 | default:
125 | }
126 | time.Sleep(100 * time.Millisecond)
127 | }
128 | }
129 |
130 | log.Printf("Unable to connect Zookeeper %s: %v\n", EmbeddedZkServer, err)
131 | return false
132 | }
133 |
134 | func startEmbeddedZookeeper(scriptRelativeDirPath string) error {
135 | err := runCmd(scriptRelativeDirPath, "start.sh")
136 | if err != nil {
137 | return err
138 | } else if !isEmbeddedZookeeperStarted(5 * time.Second) {
139 | return errors.New("embedded zk is not started")
140 | }
141 | return nil
142 | }
143 |
144 | func stopEmbeddedZookeeper(scriptRelativeDirPath string) error {
145 | return runCmd(scriptRelativeDirPath, "stop.sh")
146 | }
147 |
148 | func runCmd(scriptRelativeDirPath, scriptFileName string) error {
149 | cmd := exec.Cmd{
150 | Path: "/bin/bash",
151 | Args: []string{"/bin/bash", scriptFileName},
152 | Dir: scriptRelativeDirPath,
153 | }
154 |
155 | cmdReader, err := cmd.StdoutPipe()
156 | if err != nil {
157 | log.Println("Error creating StdoutPipe: ", err)
158 | return err
159 | }
160 |
161 | err = cmd.Start()
162 | if err != nil {
163 | return err
164 | }
165 |
166 | in := bufio.NewScanner(cmdReader)
167 | go func() {
168 | for in.Scan() {
169 | log.Println(in.Text())
170 | }
171 | }()
172 |
173 | time.Sleep(5 * time.Second) // wait some time as cmd.Start doesn't wait
174 | return nil
175 | }
176 |
--------------------------------------------------------------------------------