├── .github └── workflows │ ├── lint.yml │ └── tests.yml ├── .gitignore ├── .golangci.yml ├── LICENSE ├── README.md ├── checked_node.go ├── checked_node_test.go ├── cluster.go ├── cluster_opts.go ├── cluster_test.go ├── e2e_test.go ├── error.go ├── example_test.go ├── go.mod ├── go.sum ├── mocks_test.go ├── node.go ├── node_checker.go ├── node_discoverer.go ├── node_discoverer_test.go ├── node_picker.go ├── node_picker_test.go ├── notify.go ├── sql.go └── trace.go /.github/workflows/lint.yml: -------------------------------------------------------------------------------- 1 | name: lint 2 | on: 3 | push: 4 | tags: 5 | - v* 6 | branches: 7 | - master 8 | pull_request: 9 | workflow_dispatch: 10 | jobs: 11 | golangci: 12 | name: golangci-lint 13 | runs-on: ubuntu-latest 14 | steps: 15 | - uses: actions/checkout@v2 16 | - name: golangci-lint 17 | uses: golangci/golangci-lint-action@v6 18 | with: 19 | version: v1.62.0 20 | -------------------------------------------------------------------------------- /.github/workflows/tests.yml: -------------------------------------------------------------------------------- 1 | name: tests 2 | on: 3 | push: 4 | tags: 5 | - v* 6 | branches: 7 | - master 8 | pull_request: 9 | workflow_dispatch: 10 | jobs: 11 | test: 12 | strategy: 13 | matrix: 14 | go-version: [1.22.x, 1.23.x] 15 | os: [ubuntu-latest, macos-latest, windows-latest] 16 | env: 17 | OS: ${{ matrix.os }} 18 | GO: ${{ matrix.go-version }} 19 | runs-on: ${{ matrix.os }} 20 | steps: 21 | - name: Install Go 22 | uses: actions/setup-go@v5 23 | with: 24 | go-version: ${{ matrix.go-version }} 25 | - name: Checkout code 26 | uses: actions/checkout@v4 27 | - name: Test 28 | run: go test -race ./... -coverprofile=coverage.txt -covermode=atomic 29 | - name: Upload coverage to Codecov 30 | uses: codecov/codecov-action@v1 31 | if: matrix.os == 'ubuntu-latest' && matrix.go-version == '1.23.x' 32 | with: 33 | file: ./coverage.txt 34 | flags: unittests 35 | env_vars: OS,GO 36 | name: codecov 37 | fail_ci_if_error: true 38 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Binaries for programs and plugins 2 | *.exe 3 | *.exe~ 4 | *.dll 5 | *.so 6 | *.dylib 7 | 8 | # Test binary, built with `go test -c` 9 | *.test 10 | 11 | # Output of the go coverage tool, specifically when used with LiteIDE 12 | *.out 13 | 14 | # Dependency directories (remove the comment below to include it) 15 | # vendor/ 16 | 17 | # IDEs 18 | .idea 19 | -------------------------------------------------------------------------------- /.golangci.yml: -------------------------------------------------------------------------------- 1 | # options for analysis running 2 | run: 3 | # default concurrency is a available CPU number 4 | concurrency: 4 5 | 6 | # timeout for analysis, e.g. 30s, 5m, default is 1m 7 | deadline: 5m 8 | 9 | # exit code when at least one issue was found, default is 1 10 | issues-exit-code: 1 11 | 12 | # include test files or not, default is true 13 | tests: true 14 | 15 | # list of build tags, all linters use it. Default is empty list. 16 | #build-tags: 17 | # - mytag 18 | 19 | # which dirs to skip: they won't be analyzed; 20 | # can use regexp here: generated.*, regexp is applied on full path; 21 | # default value is empty list, but next dirs are always skipped independently 22 | # from this option's value: 23 | # vendor$, third_party$, testdata$, examples$, Godeps$, builtin$ 24 | # skip-dirs: 25 | 26 | # which files to skip: they will be analyzed, but issues from them 27 | # won't be reported. Default value is empty list, but there is 28 | # no need to include all autogenerated files, we confidently recognize 29 | # autogenerated files. If it's not please let us know. 30 | # skip-files: 31 | # - ".*\\.my\\.go$" 32 | # - lib/bad.go 33 | 34 | 35 | # output configuration options 36 | output: 37 | # colored-line-number|line-number|json|tab|checkstyle, default is "colored-line-number" 38 | format: colored-line-number 39 | 40 | # print lines of code with issue, default is true 41 | print-issued-lines: true 42 | 43 | # print linter name in the end of issue text, default is true 44 | print-linter-name: true 45 | 46 | 47 | # all available settings of specific linters 48 | linters-settings: 49 | errcheck: 50 | # report about not checking of errors in type assertions: `a := b.(MyStruct)`; 51 | # default is false: such cases aren't reported by default. 52 | check-type-assertions: false 53 | 54 | # report about assignment of errors to blank identifier: `num, _ := strconv.Atoi(numStr)`; 55 | # default is false: such cases aren't reported by default. 56 | check-blank: false 57 | govet: 58 | # report about shadowed variables 59 | check-shadowing: true 60 | golint: 61 | # minimal confidence for issues, default is 0.8 62 | min-confidence: 0.8 63 | gofmt: 64 | # simplify code: gofmt with `-s` option, true by default 65 | simplify: true 66 | goimports: 67 | # put imports beginning with prefix after 3rd-party packages; 68 | # it's a comma-separated list of prefixes 69 | local-prefixes: golang.yandex 70 | goconst: 71 | # minimal length of string constant, 3 by default 72 | min-len: 2 73 | # minimal occurrences count to trigger, 3 by default 74 | min-occurrences: 2 75 | maligned: 76 | # print struct with more effective memory layout or not, false by default 77 | suggest-new: true 78 | misspell: 79 | # Correct spellings using locale preferences for US or UK. 80 | # Default is to use a neutral variety of English. 81 | # Setting locale to US will correct the British spelling of 'colour' to 'color'. 82 | locale: US 83 | unused: 84 | # treat code as a program (not a library) and report unused exported identifiers; default is false. 85 | # XXX: if you enable this setting, unused will report a lot of false-positives in text editors: 86 | # if it's called for subdir of a project it can't find funcs usages. All text editor integrations 87 | # with golangci-lint call it on a directory with the changed file. 88 | check-exported: false 89 | unparam: 90 | # call graph construction algorithm (cha, rta). In general, use cha for libraries, 91 | # and rta for programs with main packages. Default is cha. 92 | algo: cha 93 | 94 | # Inspect exported functions, default is false. Set to true if no external program/library imports your code. 95 | # XXX: if you enable this setting, unparam will report a lot of false-positives in text editors: 96 | # if it's called for subdir of a project it can't find external interfaces. All text editor integrations 97 | # with golangci-lint call it on a directory with the changed file. 98 | check-exported: false 99 | 100 | linters: 101 | disable-all: true 102 | enable: 103 | - errcheck 104 | - goconst 105 | - gofmt # On why gofmt when goimports is enabled - https://github.com/golang/go/issues/21476 106 | - goimports 107 | - gosimple 108 | - govet 109 | - ineffassign 110 | - misspell 111 | - revive 112 | - staticcheck 113 | - typecheck 114 | - unconvert 115 | - unparam 116 | - unused 117 | 118 | issues: 119 | # List of regexps of issue texts to exclude, empty list by default. 120 | # But independently from this option we use default exclude patterns, 121 | # it can be disabled by `exclude-use-default: false`. To list all 122 | # excluded by default patterns execute `golangci-lint run --help` 123 | # exclude: 124 | 125 | # Independently from option `exclude` we use default exclude patterns, 126 | # it can be disabled by this option. To list all 127 | # excluded by default patterns execute `golangci-lint run --help`. 128 | # Default value for this option is true. 129 | exclude-use-default: false 130 | 131 | # Maximum issues count per one linter. Set to 0 to disable. Default is 50. 132 | max-per-linter: 0 133 | 134 | # Maximum count of issues with the same text. Set to 0 to disable. Default is 3. 135 | max-same-issues: 0 136 | 137 | # Show only new issues: if there are unstaged changes or untracked files, 138 | # only those changes are analyzed, else only changes in HEAD~ are analyzed. 139 | # It's a super-useful option for integration of golangci-lint into existing 140 | # large codebase. It's not practical to fix all existing issues at the moment 141 | # of integration: much better don't allow issues in new code. 142 | # Default is false. 143 | new: false 144 | 145 | # Show only new issues created after git revision `REV` 146 | # new-from-rev: REV 147 | 148 | # Show only new issues created in git patch with set file path. 149 | # new-from-patch: path/to/patch/file 150 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Apache License 2 | Version 2.0, January 2004 3 | http://www.apache.org/licenses/ 4 | 5 | TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 6 | 7 | 1. Definitions. 8 | 9 | "License" shall mean the terms and conditions for use, reproduction, 10 | and distribution as defined by Sections 1 through 9 of this document. 11 | 12 | "Licensor" shall mean the copyright owner or entity authorized by 13 | the copyright owner that is granting the License. 14 | 15 | "Legal Entity" shall mean the union of the acting entity and all 16 | other entities that control, are controlled by, or are under common 17 | control with that entity. For the purposes of this definition, 18 | "control" means (i) the power, direct or indirect, to cause the 19 | direction or management of such entity, whether by contract or 20 | otherwise, or (ii) ownership of fifty percent (50%) or more of the 21 | outstanding shares, or (iii) beneficial ownership of such entity. 22 | 23 | "You" (or "Your") shall mean an individual or Legal Entity 24 | exercising permissions granted by this License. 25 | 26 | "Source" form shall mean the preferred form for making modifications, 27 | including but not limited to software source code, documentation 28 | source, and configuration files. 29 | 30 | "Object" form shall mean any form resulting from mechanical 31 | transformation or translation of a Source form, including but 32 | not limited to compiled object code, generated documentation, 33 | and conversions to other media types. 34 | 35 | "Work" shall mean the work of authorship, whether in Source or 36 | Object form, made available under the License, as indicated by a 37 | copyright notice that is included in or attached to the work 38 | (an example is provided in the Appendix below). 39 | 40 | "Derivative Works" shall mean any work, whether in Source or Object 41 | form, that is based on (or derived from) the Work and for which the 42 | editorial revisions, annotations, elaborations, or other modifications 43 | represent, as a whole, an original work of authorship. For the purposes 44 | of this License, Derivative Works shall not include works that remain 45 | separable from, or merely link (or bind by name) to the interfaces of, 46 | the Work and Derivative Works thereof. 47 | 48 | "Contribution" shall mean any work of authorship, including 49 | the original version of the Work and any modifications or additions 50 | to that Work or Derivative Works thereof, that is intentionally 51 | submitted to Licensor for inclusion in the Work by the copyright owner 52 | or by an individual or Legal Entity authorized to submit on behalf of 53 | the copyright owner. For the purposes of this definition, "submitted" 54 | means any form of electronic, verbal, or written communication sent 55 | to the Licensor or its representatives, including but not limited to 56 | communication on electronic mailing lists, source code control systems, 57 | and issue tracking systems that are managed by, or on behalf of, the 58 | Licensor for the purpose of discussing and improving the Work, but 59 | excluding communication that is conspicuously marked or otherwise 60 | designated in writing by the copyright owner as "Not a Contribution." 61 | 62 | "Contributor" shall mean Licensor and any individual or Legal Entity 63 | on behalf of whom a Contribution has been received by Licensor and 64 | subsequently incorporated within the Work. 65 | 66 | 2. Grant of Copyright License. Subject to the terms and conditions of 67 | this License, each Contributor hereby grants to You a perpetual, 68 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 69 | copyright license to reproduce, prepare Derivative Works of, 70 | publicly display, publicly perform, sublicense, and distribute the 71 | Work and such Derivative Works in Source or Object form. 72 | 73 | 3. Grant of Patent License. Subject to the terms and conditions of 74 | this License, each Contributor hereby grants to You a perpetual, 75 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 76 | (except as stated in this section) patent license to make, have made, 77 | use, offer to sell, sell, import, and otherwise transfer the Work, 78 | where such license applies only to those patent claims licensable 79 | by such Contributor that are necessarily infringed by their 80 | Contribution(s) alone or by combination of their Contribution(s) 81 | with the Work to which such Contribution(s) was submitted. If You 82 | institute patent litigation against any entity (including a 83 | cross-claim or counterclaim in a lawsuit) alleging that the Work 84 | or a Contribution incorporated within the Work constitutes direct 85 | or contributory patent infringement, then any patent licenses 86 | granted to You under this License for that Work shall terminate 87 | as of the date such litigation is filed. 88 | 89 | 4. Redistribution. You may reproduce and distribute copies of the 90 | Work or Derivative Works thereof in any medium, with or without 91 | modifications, and in Source or Object form, provided that You 92 | meet the following conditions: 93 | 94 | (a) You must give any other recipients of the Work or 95 | Derivative Works a copy of this License; and 96 | 97 | (b) You must cause any modified files to carry prominent notices 98 | stating that You changed the files; and 99 | 100 | (c) You must retain, in the Source form of any Derivative Works 101 | that You distribute, all copyright, patent, trademark, and 102 | attribution notices from the Source form of the Work, 103 | excluding those notices that do not pertain to any part of 104 | the Derivative Works; and 105 | 106 | (d) If the Work includes a "NOTICE" text file as part of its 107 | distribution, then any Derivative Works that You distribute must 108 | include a readable copy of the attribution notices contained 109 | within such NOTICE file, excluding those notices that do not 110 | pertain to any part of the Derivative Works, in at least one 111 | of the following places: within a NOTICE text file distributed 112 | as part of the Derivative Works; within the Source form or 113 | documentation, if provided along with the Derivative Works; or, 114 | within a display generated by the Derivative Works, if and 115 | wherever such third-party notices normally appear. The contents 116 | of the NOTICE file are for informational purposes only and 117 | do not modify the License. You may add Your own attribution 118 | notices within Derivative Works that You distribute, alongside 119 | or as an addendum to the NOTICE text from the Work, provided 120 | that such additional attribution notices cannot be construed 121 | as modifying the License. 122 | 123 | You may add Your own copyright statement to Your modifications and 124 | may provide additional or different license terms and conditions 125 | for use, reproduction, or distribution of Your modifications, or 126 | for any such Derivative Works as a whole, provided Your use, 127 | reproduction, and distribution of the Work otherwise complies with 128 | the conditions stated in this License. 129 | 130 | 5. Submission of Contributions. Unless You explicitly state otherwise, 131 | any Contribution intentionally submitted for inclusion in the Work 132 | by You to the Licensor shall be under the terms and conditions of 133 | this License, without any additional terms or conditions. 134 | Notwithstanding the above, nothing herein shall supersede or modify 135 | the terms of any separate license agreement you may have executed 136 | with Licensor regarding such Contributions. 137 | 138 | 6. Trademarks. This License does not grant permission to use the trade 139 | names, trademarks, service marks, or product names of the Licensor, 140 | except as required for reasonable and customary use in describing the 141 | origin of the Work and reproducing the content of the NOTICE file. 142 | 143 | 7. Disclaimer of Warranty. Unless required by applicable law or 144 | agreed to in writing, Licensor provides the Work (and each 145 | Contributor provides its Contributions) on an "AS IS" BASIS, 146 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or 147 | implied, including, without limitation, any warranties or conditions 148 | of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A 149 | PARTICULAR PURPOSE. You are solely responsible for determining the 150 | appropriateness of using or redistributing the Work and assume any 151 | risks associated with Your exercise of permissions under this License. 152 | 153 | 8. Limitation of Liability. In no event and under no legal theory, 154 | whether in tort (including negligence), contract, or otherwise, 155 | unless required by applicable law (such as deliberate and grossly 156 | negligent acts) or agreed to in writing, shall any Contributor be 157 | liable to You for damages, including any direct, indirect, special, 158 | incidental, or consequential damages of any character arising as a 159 | result of this License or out of the use or inability to use the 160 | Work (including but not limited to damages for loss of goodwill, 161 | work stoppage, computer failure or malfunction, or any and all 162 | other commercial damages or losses), even if such Contributor 163 | has been advised of the possibility of such damages. 164 | 165 | 9. Accepting Warranty or Additional Liability. While redistributing 166 | the Work or Derivative Works thereof, You may choose to offer, 167 | and charge a fee for, acceptance of support, warranty, indemnity, 168 | or other liability obligations and/or rights consistent with this 169 | License. However, in accepting such obligations, You may act only 170 | on Your own behalf and on Your sole responsibility, not on behalf 171 | of any other Contributor, and only if You agree to indemnify, 172 | defend, and hold each Contributor harmless for any liability 173 | incurred by, or claims asserted against, such Contributor by reason 174 | of your accepting any such warranty or additional liability. 175 | 176 | END OF TERMS AND CONDITIONS 177 | 178 | APPENDIX: How to apply the Apache License to your work. 179 | 180 | To apply the Apache License to your work, attach the following 181 | boilerplate notice, with the fields enclosed by brackets "[]" 182 | replaced with your own identifying information. (Don't include 183 | the brackets!) The text should be enclosed in the appropriate 184 | comment syntax for the file format. We also recommend that a 185 | file or class name and description of purpose be included on the 186 | same "printed page" as the copyright notice for easier 187 | identification within third-party archives. 188 | 189 | Copyright [yyyy] [name of copyright owner] 190 | 191 | Licensed under the Apache License, Version 2.0 (the "License"); 192 | you may not use this file except in compliance with the License. 193 | You may obtain a copy of the License at 194 | 195 | http://www.apache.org/licenses/LICENSE-2.0 196 | 197 | Unless required by applicable law or agreed to in writing, software 198 | distributed under the License is distributed on an "AS IS" BASIS, 199 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 200 | See the License for the specific language governing permissions and 201 | limitations under the License. 202 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # hasql 2 | 3 | [![PkgGoDev](https://pkg.go.dev/badge/golang.yandex/hasql)](https://pkg.go.dev/golang.yandex/hasql) 4 | [![GoDoc](https://godoc.org/golang.yandex/hasql?status.svg)](https://godoc.org/golang.yandex/hasql) 5 | ![tests](https://github.com/yandex/go-hasql/workflows/tests/badge.svg?branch=master) 6 | ![lint](https://github.com/yandex/go-hasql/workflows/lint/badge.svg?branch=master) 7 | [![Go Report Card](https://goreportcard.com/badge/golang.yandex/hasql)](https://goreportcard.com/report/golang.yandex/hasql) 8 | [![codecov](https://codecov.io/gh/yandex/go-hasql/branch/master/graph/badge.svg)](https://codecov.io/gh/yandex/go-hasql) 9 | 10 | `hasql` provides simple and reliable way to access high-availability database setups with multiple hosts. 11 | 12 | ## Status 13 | `hasql` is production-ready and is actively used inside Yandex' production environment. 14 | 15 | ## Prerequisites 16 | 17 | - **[Go](https://golang.org)**: any one of the **two latest major** [releases](https://golang.org/doc/devel/release.html). 18 | 19 | ## Installation 20 | 21 | With [Go module](https://github.com/golang/go/wiki/Modules) support, simply add the following import 22 | 23 | ```go 24 | import "golang.yandex/hasql/v2" 25 | ``` 26 | 27 | to your code, and then `go [build|run|test]` will automatically fetch the 28 | necessary dependencies. 29 | 30 | Otherwise, to install the `hasql` package, run the following command: 31 | 32 | ```console 33 | $ go get -u golang.yandex/hasql/v2 34 | ``` 35 | 36 | ## How does it work 37 | `hasql` operates using standard `database/sql` connection pool objects. User creates `*sql.DB`-compatible objects for each node of database cluster and passes them to constructor. Library keeps up to date information on state of each node by 'pinging' them periodically. User is provided with a set of interfaces to retrieve database node object suitable for required operation. 38 | 39 | ```go 40 | dbFoo, _ := sql.Open("pgx", "host=foo") 41 | dbBar, _ := sql.Open("pgx", "host=bar") 42 | 43 | discoverer := NewStaticNodeDiscoverer( 44 | NewNode("foo", dbFoo), 45 | NewNode("bar", dbBar), 46 | ) 47 | 48 | cl, err := hasql.NewCluster(discoverer, hasql.PostgreSQLChecker) 49 | if err != nil { ... } 50 | 51 | node := cl.Primary() 52 | if node == nil { 53 | err := cl.Err() // most recent errors for all nodes in the cluster 54 | } 55 | 56 | fmt.Printf("got node %s\n", node) 57 | 58 | ctx, cancel := context.WithTimeout(context.Background(), time.Second) 59 | defer cancel() 60 | 61 | if err = node.DB().PingContext(ctx); err != nil { ... } 62 | ``` 63 | 64 | `hasql` does not configure provided connection pools in any way. It is user's job to set them up properly. Library does handle their lifetime though - pools are closed when `Cluster` object is closed. 65 | 66 | ### Concepts and entities 67 | 68 | **Cluster** is a set of `database/sql`-compatible nodes that tracks their lifespan and provides access to each individual nodes. 69 | 70 | **Node** is a single database instance in high-availability cluster. 71 | 72 | **Node discoverer** provides nodes objects to cluster. This abstraction allows user to dynamically change set of cluster nodes, for example collect nodes list via Service Discovery (etcd, Consul). 73 | 74 | **Node checker** collects information about current state of individual node in cluster, such as: cluster role, network latency, replication lag, etc. 75 | 76 | **Node picker** picks node from cluster by given criterion using predefined algorithm: random, round-robin, lowest latency, etc. 77 | 78 | ### Supported criteria 79 | _Alive primary_|_Alive standby_|_Any alive_ node or _none_ otherwise 80 | ```go 81 | node := c.Node(hasql.Alive) 82 | ``` 83 | 84 | _Alive primary_ or _none_ otherwise 85 | ```go 86 | node := c.Node(hasql.Primary) 87 | ``` 88 | 89 | _Alive standby_ or _none_ otherwise 90 | ```go 91 | node := c.Node(hasql.Standby) 92 | ``` 93 | 94 | _Alive primary_|_Alive standby_ or _none_ otherwise 95 | ```go 96 | node := c.Node(hasql.PreferPrimary) 97 | ``` 98 | 99 | _Alive standby_|_Alive primary_ or _none_ otherwise 100 | ```go 101 | node := c.Node(hasql.PreferStandby) 102 | ``` 103 | 104 | ### Node pickers 105 | When user asks `Cluster` object for a node a random one from a list of suitable nodes is returned. User can override this behavior by providing a custom node picker. 106 | 107 | Library provides a couple of predefined pickers. For example if user wants 'closest' node (with lowest latency) `LatencyNodePicker` should be used. 108 | 109 | ```go 110 | cl, err := hasql.NewCluster( 111 | hasql.NewStaticNodeDiscoverer(hasql.NewNode("foo", dbFoo), hasql.NewNode("bar", dbBar)), 112 | hasql.PostgreSQLChecker, 113 | hasql.WithNodePicker(new(hasql.LatencyNodePicker[*sql.DB])), 114 | ) 115 | ``` 116 | 117 | ## Supported databases 118 | Since library requires `Querier` interface, which describes a subset of `database/sql.DB` methods, it supports any database that has a `database/sql` driver. All it requires is a database-specific checker function that can provide node state info. 119 | 120 | Check out `node_checker.go` file for more information. 121 | 122 | ### Caveats 123 | Node's state is transient at any given time. If `Primary()` returns a node it does not mean that node is still primary when you execute statement on it. All it means is that it was primary when it was last checked. Nodes can change their state at a whim or even go offline and `hasql` can't control it in any manner. 124 | 125 | This is one of the reasons why nodes do not expose their perceived state to user. 126 | 127 | ## Extensions 128 | ### Instrumentation 129 | You can add instrumentation via `Tracer` object similar to [httptrace](https://godoc.org/net/http/httptrace) in standard library. 130 | -------------------------------------------------------------------------------- /checked_node.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright 2020 YANDEX LLC 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | */ 16 | 17 | package hasql 18 | 19 | import ( 20 | "context" 21 | "errors" 22 | "fmt" 23 | "slices" 24 | "sync" 25 | ) 26 | 27 | // CheckedNodes holds references to all available cluster nodes 28 | type CheckedNodes[T Querier] struct { 29 | discovered []*Node[T] 30 | alive []CheckedNode[T] 31 | primaries []CheckedNode[T] 32 | standbys []CheckedNode[T] 33 | err error 34 | } 35 | 36 | // Discovered returns a list of nodes discovered in cluster 37 | func (c CheckedNodes[T]) Discovered() []*Node[T] { 38 | return c.discovered 39 | } 40 | 41 | // Alive returns a list of all successfully checked nodes irregarding their cluster role 42 | func (c CheckedNodes[T]) Alive() []CheckedNode[T] { 43 | return c.alive 44 | } 45 | 46 | // Primaries returns list of all successfully checked nodes with primary role 47 | func (c CheckedNodes[T]) Primaries() []CheckedNode[T] { 48 | return c.primaries 49 | } 50 | 51 | // Standbys returns list of all successfully checked nodes with standby role 52 | func (c CheckedNodes[T]) Standbys() []CheckedNode[T] { 53 | return c.standbys 54 | } 55 | 56 | // Err holds information about cause of node check failure. 57 | func (c CheckedNodes[T]) Err() error { 58 | return c.err 59 | } 60 | 61 | // CheckedNode contains most recent state of single cluster node 62 | type CheckedNode[T Querier] struct { 63 | Node *Node[T] 64 | Info NodeInfoProvider 65 | } 66 | 67 | // checkNodes takes slice of nodes, checks them in parallel and returns the alive ones 68 | func checkNodes[T Querier](ctx context.Context, discoverer NodeDiscoverer[T], checkFn NodeChecker, compareFn func(a, b CheckedNode[T]) int, tracer Tracer[T]) CheckedNodes[T] { 69 | discoveredNodes, err := discoverer.DiscoverNodes(ctx) 70 | if err != nil { 71 | // error discovering nodes 72 | return CheckedNodes[T]{ 73 | err: fmt.Errorf("cannot discover cluster nodes: %w", err), 74 | } 75 | } 76 | 77 | var mu sync.Mutex 78 | checked := make([]CheckedNode[T], 0, len(discoveredNodes)) 79 | var errs NodeCheckErrors[T] 80 | 81 | var wg sync.WaitGroup 82 | wg.Add(len(discoveredNodes)) 83 | for _, node := range discoveredNodes { 84 | go func(node *Node[T]) { 85 | defer wg.Done() 86 | 87 | // check single node state 88 | info, err := checkFn(ctx, node.DB()) 89 | if err != nil { 90 | cerr := NodeCheckError[T]{ 91 | node: node, 92 | err: err, 93 | } 94 | 95 | // node is dead - make trace call 96 | if tracer.NodeDead != nil { 97 | tracer.NodeDead(cerr) 98 | } 99 | 100 | // store node check error 101 | mu.Lock() 102 | defer mu.Unlock() 103 | errs = append(errs, cerr) 104 | return 105 | } 106 | 107 | cn := CheckedNode[T]{ 108 | Node: node, 109 | Info: info, 110 | } 111 | 112 | // make trace call about alive node 113 | if tracer.NodeAlive != nil { 114 | tracer.NodeAlive(cn) 115 | } 116 | 117 | // store checked alive node 118 | mu.Lock() 119 | defer mu.Unlock() 120 | checked = append(checked, cn) 121 | }(node) 122 | } 123 | 124 | // wait for all nodes to be checked 125 | wg.Wait() 126 | 127 | // sort checked nodes 128 | slices.SortFunc(checked, compareFn) 129 | 130 | // split checked nodes by roles 131 | alive := make([]CheckedNode[T], 0, len(checked)) 132 | // in almost all cases there is only one primary node in cluster 133 | primaries := make([]CheckedNode[T], 0, 1) 134 | standbys := make([]CheckedNode[T], 0, len(checked)) 135 | for _, cn := range checked { 136 | switch cn.Info.Role() { 137 | case NodeRolePrimary: 138 | primaries = append(primaries, cn) 139 | alive = append(alive, cn) 140 | case NodeRoleStandby: 141 | standbys = append(standbys, cn) 142 | alive = append(alive, cn) 143 | default: 144 | // treat node with undetermined role as dead 145 | cerr := NodeCheckError[T]{ 146 | node: cn.Node, 147 | err: errors.New("cannot determine node role"), 148 | } 149 | errs = append(errs, cerr) 150 | 151 | if tracer.NodeDead != nil { 152 | tracer.NodeDead(cerr) 153 | } 154 | } 155 | } 156 | 157 | res := CheckedNodes[T]{ 158 | discovered: discoveredNodes, 159 | alive: alive, 160 | primaries: primaries, 161 | standbys: standbys, 162 | err: func() error { 163 | if len(errs) != 0 { 164 | return errs 165 | } 166 | return nil 167 | }(), 168 | } 169 | 170 | return res 171 | } 172 | 173 | // pickNodeByCriterion is a helper function to pick a single node by given criterion 174 | func pickNodeByCriterion[T Querier](nodes CheckedNodes[T], picker NodePicker[T], criterion NodeStateCriterion) *Node[T] { 175 | var subset []CheckedNode[T] 176 | 177 | switch criterion { 178 | case Alive: 179 | subset = nodes.alive 180 | case Primary: 181 | subset = nodes.primaries 182 | case Standby: 183 | subset = nodes.standbys 184 | case PreferPrimary: 185 | if subset = nodes.primaries; len(subset) == 0 { 186 | subset = nodes.standbys 187 | } 188 | case PreferStandby: 189 | if subset = nodes.standbys; len(subset) == 0 { 190 | subset = nodes.primaries 191 | } 192 | } 193 | 194 | if len(subset) == 0 { 195 | return nil 196 | } 197 | 198 | return picker.PickNode(subset).Node 199 | } 200 | -------------------------------------------------------------------------------- /checked_node_test.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright 2020 YANDEX LLC 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | */ 16 | 17 | package hasql 18 | 19 | import ( 20 | "context" 21 | "database/sql" 22 | "errors" 23 | "io" 24 | "testing" 25 | 26 | "github.com/stretchr/testify/assert" 27 | ) 28 | 29 | func TestCheckNodes(t *testing.T) { 30 | t.Run("discovery_error", func(t *testing.T) { 31 | discoverer := mockNodeDiscoverer[*sql.DB]{ 32 | err: io.EOF, 33 | } 34 | 35 | nodes := checkNodes(context.Background(), discoverer, nil, nil, Tracer[*sql.DB]{}) 36 | assert.Empty(t, nodes.discovered) 37 | assert.Empty(t, nodes.alive) 38 | assert.Empty(t, nodes.primaries) 39 | assert.Empty(t, nodes.standbys) 40 | assert.ErrorIs(t, nodes.err, io.EOF) 41 | }) 42 | 43 | t.Run("all_nodes_alive", func(t *testing.T) { 44 | node1 := &Node[*mockQuerier]{ 45 | name: "shimba", 46 | db: &mockQuerier{name: "primary"}, 47 | } 48 | node2 := &Node[*mockQuerier]{ 49 | name: "boomba", 50 | db: &mockQuerier{name: "standby1"}, 51 | } 52 | node3 := &Node[*mockQuerier]{ 53 | name: "looken", 54 | db: &mockQuerier{name: "standby2"}, 55 | } 56 | 57 | discoverer := mockNodeDiscoverer[*mockQuerier]{ 58 | nodes: []*Node[*mockQuerier]{node1, node2, node3}, 59 | } 60 | 61 | // mock node checker func 62 | checkFn := func(_ context.Context, q Querier) (NodeInfoProvider, error) { 63 | mq, ok := q.(*mockQuerier) 64 | if !ok { 65 | return NodeInfo{}, nil 66 | } 67 | 68 | switch mq.name { 69 | case node1.db.name: 70 | return NodeInfo{ClusterRole: NodeRolePrimary, NetworkLatency: 100}, nil 71 | case node2.db.name: 72 | return NodeInfo{ClusterRole: NodeRoleStandby, NetworkLatency: 50}, nil 73 | case node3.db.name: 74 | return NodeInfo{ClusterRole: NodeRoleStandby, NetworkLatency: 70}, nil 75 | default: 76 | return NodeInfo{}, nil 77 | } 78 | } 79 | 80 | var picker LatencyNodePicker[*mockQuerier] 81 | var tracer Tracer[*mockQuerier] 82 | 83 | checked := checkNodes(context.Background(), discoverer, checkFn, picker.CompareNodes, tracer) 84 | 85 | expected := CheckedNodes[*mockQuerier]{ 86 | discovered: []*Node[*mockQuerier]{node1, node2, node3}, 87 | alive: []CheckedNode[*mockQuerier]{ 88 | {Node: node2, Info: NodeInfo{ClusterRole: NodeRoleStandby, NetworkLatency: 50}}, 89 | {Node: node3, Info: NodeInfo{ClusterRole: NodeRoleStandby, NetworkLatency: 70}}, 90 | {Node: node1, Info: NodeInfo{ClusterRole: NodeRolePrimary, NetworkLatency: 100}}, 91 | }, 92 | primaries: []CheckedNode[*mockQuerier]{ 93 | {Node: node1, Info: NodeInfo{ClusterRole: NodeRolePrimary, NetworkLatency: 100}}, 94 | }, 95 | standbys: []CheckedNode[*mockQuerier]{ 96 | {Node: node2, Info: NodeInfo{ClusterRole: NodeRoleStandby, NetworkLatency: 50}}, 97 | {Node: node3, Info: NodeInfo{ClusterRole: NodeRoleStandby, NetworkLatency: 70}}, 98 | }, 99 | } 100 | 101 | assert.Equal(t, expected, checked) 102 | }) 103 | 104 | t.Run("all_nodes_dead", func(t *testing.T) { 105 | node1 := &Node[*mockQuerier]{ 106 | name: "shimba", 107 | db: &mockQuerier{name: "primary"}, 108 | } 109 | node2 := &Node[*mockQuerier]{ 110 | name: "boomba", 111 | db: &mockQuerier{name: "standby1"}, 112 | } 113 | node3 := &Node[*mockQuerier]{ 114 | name: "looken", 115 | db: &mockQuerier{name: "standby2"}, 116 | } 117 | 118 | discoverer := mockNodeDiscoverer[*mockQuerier]{ 119 | nodes: []*Node[*mockQuerier]{node1, node2, node3}, 120 | } 121 | 122 | // mock node checker func 123 | checkFn := func(_ context.Context, _ Querier) (NodeInfoProvider, error) { 124 | return nil, io.EOF 125 | } 126 | 127 | var picker LatencyNodePicker[*mockQuerier] 128 | var tracer Tracer[*mockQuerier] 129 | 130 | checked := checkNodes(context.Background(), discoverer, checkFn, picker.CompareNodes, tracer) 131 | 132 | expectedDiscovered := []*Node[*mockQuerier]{node1, node2, node3} 133 | assert.Equal(t, expectedDiscovered, checked.discovered) 134 | 135 | assert.Empty(t, checked.alive) 136 | assert.Empty(t, checked.primaries) 137 | assert.Empty(t, checked.standbys) 138 | 139 | var cerrs NodeCheckErrors[*mockQuerier] 140 | assert.ErrorAs(t, checked.err, &cerrs) 141 | assert.Len(t, cerrs, 3) 142 | for _, cerr := range cerrs { 143 | assert.ErrorIs(t, cerr, io.EOF) 144 | } 145 | }) 146 | 147 | t.Run("one_standby_is_dead", func(t *testing.T) { 148 | node1 := &Node[*mockQuerier]{ 149 | name: "shimba", 150 | db: &mockQuerier{name: "primary"}, 151 | } 152 | node2 := &Node[*mockQuerier]{ 153 | name: "boomba", 154 | db: &mockQuerier{name: "standby1"}, 155 | } 156 | node3 := &Node[*mockQuerier]{ 157 | name: "looken", 158 | db: &mockQuerier{name: "standby2"}, 159 | } 160 | 161 | discoverer := mockNodeDiscoverer[*mockQuerier]{ 162 | nodes: []*Node[*mockQuerier]{node1, node2, node3}, 163 | } 164 | 165 | // mock node checker func 166 | checkFn := func(_ context.Context, q Querier) (NodeInfoProvider, error) { 167 | mq, ok := q.(*mockQuerier) 168 | if !ok { 169 | return NodeInfo{}, nil 170 | } 171 | 172 | switch mq.name { 173 | case node1.db.name: 174 | return NodeInfo{ClusterRole: NodeRolePrimary, NetworkLatency: 100}, nil 175 | case node2.db.name: 176 | return NodeInfo{ClusterRole: NodeRoleStandby, NetworkLatency: 50}, nil 177 | case node3.db.name: 178 | return nil, io.EOF 179 | default: 180 | return NodeInfo{}, nil 181 | } 182 | } 183 | 184 | var picker LatencyNodePicker[*mockQuerier] 185 | var tracer Tracer[*mockQuerier] 186 | 187 | checked := checkNodes(context.Background(), discoverer, checkFn, picker.CompareNodes, tracer) 188 | 189 | expected := CheckedNodes[*mockQuerier]{ 190 | discovered: []*Node[*mockQuerier]{node1, node2, node3}, 191 | alive: []CheckedNode[*mockQuerier]{ 192 | {Node: node2, Info: NodeInfo{ClusterRole: NodeRoleStandby, NetworkLatency: 50}}, 193 | {Node: node1, Info: NodeInfo{ClusterRole: NodeRolePrimary, NetworkLatency: 100}}, 194 | }, 195 | primaries: []CheckedNode[*mockQuerier]{ 196 | {Node: node1, Info: NodeInfo{ClusterRole: NodeRolePrimary, NetworkLatency: 100}}, 197 | }, 198 | standbys: []CheckedNode[*mockQuerier]{ 199 | {Node: node2, Info: NodeInfo{ClusterRole: NodeRoleStandby, NetworkLatency: 50}}, 200 | }, 201 | err: NodeCheckErrors[*mockQuerier]{ 202 | {node: node3, err: io.EOF}, 203 | }, 204 | } 205 | 206 | assert.Equal(t, expected, checked) 207 | }) 208 | 209 | t.Run("primary_is_dead", func(t *testing.T) { 210 | node1 := &Node[*mockQuerier]{ 211 | name: "shimba", 212 | db: &mockQuerier{name: "primary"}, 213 | } 214 | node2 := &Node[*mockQuerier]{ 215 | name: "boomba", 216 | db: &mockQuerier{name: "standby1"}, 217 | } 218 | node3 := &Node[*mockQuerier]{ 219 | name: "looken", 220 | db: &mockQuerier{name: "standby2"}, 221 | } 222 | 223 | discoverer := mockNodeDiscoverer[*mockQuerier]{ 224 | nodes: []*Node[*mockQuerier]{node1, node2, node3}, 225 | } 226 | 227 | // mock node checker func 228 | checkFn := func(_ context.Context, q Querier) (NodeInfoProvider, error) { 229 | mq, ok := q.(*mockQuerier) 230 | if !ok { 231 | return NodeInfo{}, nil 232 | } 233 | 234 | switch mq.name { 235 | case node1.db.name: 236 | return nil, io.EOF 237 | case node2.db.name: 238 | return NodeInfo{ClusterRole: NodeRoleStandby, NetworkLatency: 70}, nil 239 | case node3.db.name: 240 | return NodeInfo{ClusterRole: NodeRoleStandby, NetworkLatency: 50}, nil 241 | default: 242 | return NodeInfo{}, nil 243 | } 244 | } 245 | 246 | var picker LatencyNodePicker[*mockQuerier] 247 | var tracer Tracer[*mockQuerier] 248 | 249 | checked := checkNodes(context.Background(), discoverer, checkFn, picker.CompareNodes, tracer) 250 | 251 | expected := CheckedNodes[*mockQuerier]{ 252 | discovered: []*Node[*mockQuerier]{node1, node2, node3}, 253 | alive: []CheckedNode[*mockQuerier]{ 254 | {Node: node3, Info: NodeInfo{ClusterRole: NodeRoleStandby, NetworkLatency: 50}}, 255 | {Node: node2, Info: NodeInfo{ClusterRole: NodeRoleStandby, NetworkLatency: 70}}, 256 | }, 257 | primaries: []CheckedNode[*mockQuerier]{}, 258 | standbys: []CheckedNode[*mockQuerier]{ 259 | {Node: node3, Info: NodeInfo{ClusterRole: NodeRoleStandby, NetworkLatency: 50}}, 260 | {Node: node2, Info: NodeInfo{ClusterRole: NodeRoleStandby, NetworkLatency: 70}}, 261 | }, 262 | err: NodeCheckErrors[*mockQuerier]{ 263 | {node: node1, err: io.EOF}, 264 | }, 265 | } 266 | 267 | assert.Equal(t, expected, checked) 268 | }) 269 | 270 | t.Run("node_with_unknown_role", func(t *testing.T) { 271 | node1 := &Node[*mockQuerier]{ 272 | name: "shimba", 273 | db: &mockQuerier{name: "unknown"}, 274 | } 275 | node2 := &Node[*mockQuerier]{ 276 | name: "boomba", 277 | db: &mockQuerier{name: "primary"}, 278 | } 279 | node3 := &Node[*mockQuerier]{ 280 | name: "looken", 281 | db: &mockQuerier{name: "standby2"}, 282 | } 283 | 284 | discoverer := mockNodeDiscoverer[*mockQuerier]{ 285 | nodes: []*Node[*mockQuerier]{node1, node2, node3}, 286 | } 287 | 288 | // mock node checker func 289 | checkFn := func(_ context.Context, q Querier) (NodeInfoProvider, error) { 290 | mq, ok := q.(*mockQuerier) 291 | if !ok { 292 | return NodeInfo{}, nil 293 | } 294 | 295 | switch mq.name { 296 | case node1.db.name: 297 | return NodeInfo{}, nil 298 | case node2.db.name: 299 | return NodeInfo{ClusterRole: NodeRolePrimary, NetworkLatency: 20}, nil 300 | case node3.db.name: 301 | return NodeInfo{ClusterRole: NodeRoleStandby, NetworkLatency: 50}, nil 302 | default: 303 | return NodeInfo{}, nil 304 | } 305 | } 306 | 307 | var picker LatencyNodePicker[*mockQuerier] 308 | var tracer Tracer[*mockQuerier] 309 | 310 | checked := checkNodes(context.Background(), discoverer, checkFn, picker.CompareNodes, tracer) 311 | 312 | expected := CheckedNodes[*mockQuerier]{ 313 | discovered: []*Node[*mockQuerier]{node1, node2, node3}, 314 | alive: []CheckedNode[*mockQuerier]{ 315 | {Node: node2, Info: NodeInfo{ClusterRole: NodeRolePrimary, NetworkLatency: 20}}, 316 | {Node: node3, Info: NodeInfo{ClusterRole: NodeRoleStandby, NetworkLatency: 50}}, 317 | }, 318 | primaries: []CheckedNode[*mockQuerier]{ 319 | {Node: node2, Info: NodeInfo{ClusterRole: NodeRolePrimary, NetworkLatency: 20}}, 320 | }, 321 | standbys: []CheckedNode[*mockQuerier]{ 322 | {Node: node3, Info: NodeInfo{ClusterRole: NodeRoleStandby, NetworkLatency: 50}}, 323 | }, 324 | err: NodeCheckErrors[*mockQuerier]{ 325 | {node: node1, err: errors.New("cannot determine node role")}, 326 | }, 327 | } 328 | 329 | assert.Equal(t, expected, checked) 330 | }) 331 | } 332 | 333 | func TestPickNodeByCriterion(t *testing.T) { 334 | t.Run("no_nodes", func(t *testing.T) { 335 | nodes := CheckedNodes[*sql.DB]{} 336 | picker := new(RandomNodePicker[*sql.DB]) 337 | 338 | // all criteria must return nil node 339 | for i := Alive; i < maxNodeCriterion; i++ { 340 | node := pickNodeByCriterion(nodes, picker, i) 341 | assert.Nil(t, node) 342 | } 343 | }) 344 | 345 | t.Run("alive", func(t *testing.T) { 346 | node := &Node[*sql.DB]{ 347 | name: "shimba", 348 | db: new(sql.DB), 349 | } 350 | 351 | nodes := CheckedNodes[*sql.DB]{ 352 | alive: []CheckedNode[*sql.DB]{ 353 | { 354 | Node: node, 355 | Info: NodeInfo{ 356 | ClusterRole: NodeRolePrimary, 357 | }, 358 | }, 359 | }, 360 | } 361 | picker := new(RandomNodePicker[*sql.DB]) 362 | 363 | assert.Equal(t, node, pickNodeByCriterion(nodes, picker, Alive)) 364 | }) 365 | 366 | t.Run("primary", func(t *testing.T) { 367 | node := &Node[*sql.DB]{ 368 | name: "shimba", 369 | db: new(sql.DB), 370 | } 371 | 372 | nodes := CheckedNodes[*sql.DB]{ 373 | primaries: []CheckedNode[*sql.DB]{ 374 | { 375 | Node: node, 376 | Info: NodeInfo{ 377 | ClusterRole: NodeRolePrimary, 378 | }, 379 | }, 380 | }, 381 | } 382 | picker := new(RandomNodePicker[*sql.DB]) 383 | 384 | assert.Equal(t, node, pickNodeByCriterion(nodes, picker, Primary)) 385 | // we will return node on Prefer* criteria also 386 | assert.Equal(t, node, pickNodeByCriterion(nodes, picker, PreferPrimary)) 387 | assert.Equal(t, node, pickNodeByCriterion(nodes, picker, PreferStandby)) 388 | }) 389 | 390 | t.Run("standby", func(t *testing.T) { 391 | node := &Node[*sql.DB]{ 392 | name: "shimba", 393 | db: new(sql.DB), 394 | } 395 | 396 | nodes := CheckedNodes[*sql.DB]{ 397 | standbys: []CheckedNode[*sql.DB]{ 398 | { 399 | Node: node, 400 | Info: NodeInfo{ 401 | ClusterRole: NodeRoleStandby, 402 | }, 403 | }, 404 | }, 405 | } 406 | picker := new(RandomNodePicker[*sql.DB]) 407 | 408 | assert.Equal(t, node, pickNodeByCriterion(nodes, picker, Standby)) 409 | // we will return node on Prefer* criteria also 410 | assert.Equal(t, node, pickNodeByCriterion(nodes, picker, PreferPrimary)) 411 | assert.Equal(t, node, pickNodeByCriterion(nodes, picker, PreferStandby)) 412 | }) 413 | 414 | t.Run("prefer_primary", func(t *testing.T) { 415 | node := &Node[*sql.DB]{ 416 | name: "shimba", 417 | db: new(sql.DB), 418 | } 419 | 420 | // must pick from primaries 421 | nodes := CheckedNodes[*sql.DB]{ 422 | primaries: []CheckedNode[*sql.DB]{ 423 | { 424 | Node: node, 425 | Info: NodeInfo{ 426 | ClusterRole: NodeRolePrimary, 427 | }, 428 | }, 429 | }, 430 | } 431 | picker := new(RandomNodePicker[*sql.DB]) 432 | 433 | assert.Equal(t, node, pickNodeByCriterion(nodes, picker, PreferPrimary)) 434 | 435 | // must pick from standbys 436 | nodes = CheckedNodes[*sql.DB]{ 437 | standbys: []CheckedNode[*sql.DB]{ 438 | { 439 | Node: node, 440 | Info: NodeInfo{ 441 | ClusterRole: NodeRoleStandby, 442 | }, 443 | }, 444 | }, 445 | } 446 | assert.Equal(t, node, pickNodeByCriterion(nodes, picker, PreferPrimary)) 447 | }) 448 | 449 | t.Run("prefer_standby", func(t *testing.T) { 450 | node := &Node[*sql.DB]{ 451 | name: "shimba", 452 | db: new(sql.DB), 453 | } 454 | 455 | // must pick from standbys 456 | nodes := CheckedNodes[*sql.DB]{ 457 | standbys: []CheckedNode[*sql.DB]{ 458 | { 459 | Node: node, 460 | Info: NodeInfo{ 461 | ClusterRole: NodeRoleStandby, 462 | }, 463 | }, 464 | }, 465 | } 466 | picker := new(RandomNodePicker[*sql.DB]) 467 | 468 | assert.Equal(t, node, pickNodeByCriterion(nodes, picker, PreferStandby)) 469 | 470 | // must pick from primaries 471 | nodes = CheckedNodes[*sql.DB]{ 472 | primaries: []CheckedNode[*sql.DB]{ 473 | { 474 | Node: node, 475 | Info: NodeInfo{ 476 | ClusterRole: NodeRolePrimary, 477 | }, 478 | }, 479 | }, 480 | } 481 | assert.Equal(t, node, pickNodeByCriterion(nodes, picker, PreferStandby)) 482 | }) 483 | } 484 | -------------------------------------------------------------------------------- /cluster.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright 2020 YANDEX LLC 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | */ 16 | 17 | // Package hasql provides simple and reliable way to access high-availability database setups with multiple hosts. 18 | package hasql 19 | 20 | import ( 21 | "context" 22 | "errors" 23 | "io" 24 | "sync" 25 | "sync/atomic" 26 | "time" 27 | ) 28 | 29 | // Cluster consists of number of 'nodes' of a single SQL database. 30 | type Cluster[T Querier] struct { 31 | // configuration 32 | updateInterval time.Duration 33 | updateTimeout time.Duration 34 | discoverer NodeDiscoverer[T] 35 | checker NodeChecker 36 | picker NodePicker[T] 37 | tracer Tracer[T] 38 | 39 | // status 40 | checkedNodes atomic.Value 41 | stop context.CancelFunc 42 | 43 | // broadcast 44 | subscribersMu sync.Mutex 45 | subscribers []updateSubscriber[T] 46 | } 47 | 48 | // NewCluster returns object representing a single 'cluster' of SQL databases 49 | func NewCluster[T Querier](discoverer NodeDiscoverer[T], checker NodeChecker, opts ...ClusterOpt[T]) (*Cluster[T], error) { 50 | if discoverer == nil { 51 | return nil, errors.New("node discoverer required") 52 | } 53 | 54 | // prepare internal 'stop' context 55 | ctx, stopFn := context.WithCancel(context.Background()) 56 | 57 | cl := &Cluster[T]{ 58 | updateInterval: 5 * time.Second, 59 | updateTimeout: time.Second, 60 | discoverer: discoverer, 61 | checker: checker, 62 | picker: new(RandomNodePicker[T]), 63 | 64 | stop: stopFn, 65 | } 66 | 67 | // apply options 68 | for _, opt := range opts { 69 | opt(cl) 70 | } 71 | 72 | // Store initial nodes state 73 | cl.checkedNodes.Store(CheckedNodes[T]{}) 74 | 75 | // Start update routine 76 | go cl.backgroundNodesUpdate(ctx) 77 | return cl, nil 78 | } 79 | 80 | // Close stops node updates. 81 | // Close function must be called when cluster is not needed anymore. 82 | // It returns combined error if multiple nodes returned errors 83 | func (cl *Cluster[T]) Close() (err error) { 84 | cl.stop() 85 | 86 | // close all nodes underlying connection pools 87 | discovered := cl.checkedNodes.Load().(CheckedNodes[T]).discovered 88 | for _, node := range discovered { 89 | if closer, ok := any(node.DB()).(io.Closer); ok { 90 | err = errors.Join(err, closer.Close()) 91 | } 92 | } 93 | 94 | // discard any collected state of nodes 95 | cl.checkedNodes.Store(CheckedNodes[T]{}) 96 | 97 | return err 98 | } 99 | 100 | // Err returns cause of nodes most recent check failures. 101 | // In most cases error is a list of errors of type CheckNodeErrors, original errors 102 | // could be extracted using `errors.As`. 103 | // Example: 104 | // 105 | // var cerrs NodeCheckErrors 106 | // if errors.As(err, &cerrs) { 107 | // for _, cerr := range cerrs { 108 | // fmt.Printf("node: %s, err: %s\n", cerr.Node(), cerr.Err()) 109 | // } 110 | // } 111 | func (cl *Cluster[T]) Err() error { 112 | return cl.checkedNodes.Load().(CheckedNodes[T]).Err() 113 | } 114 | 115 | // Node returns cluster node with specified status 116 | func (cl *Cluster[T]) Node(criterion NodeStateCriterion) *Node[T] { 117 | return pickNodeByCriterion(cl.checkedNodes.Load().(CheckedNodes[T]), cl.picker, criterion) 118 | } 119 | 120 | // WaitForNode with specified status to appear or until context is canceled 121 | func (cl *Cluster[T]) WaitForNode(ctx context.Context, criterion NodeStateCriterion) (*Node[T], error) { 122 | // Node already exists? 123 | node := cl.Node(criterion) 124 | if node != nil { 125 | return node, nil 126 | } 127 | 128 | ch := cl.addUpdateSubscriber(criterion) 129 | 130 | // Node might have appeared while we were adding waiter, recheck 131 | node = cl.Node(criterion) 132 | if node != nil { 133 | return node, nil 134 | } 135 | 136 | // If channel is unbuffered and we are right here when nodes are updated, 137 | // the update code won't be able to write into channel and will 'forget' it. 138 | // Then we will report nil to the caller, either because update code 139 | // closes channel or because context is canceled. 140 | // 141 | // In both cases its not what user wants. 142 | // 143 | // We can solve it by doing cl.Node(ns) if/when we are about to return nil. 144 | // But if another update runs between channel read and cl.Node(ns) AND no 145 | // nodes have requested status, we will still return nil. 146 | // 147 | // Also code becomes more complex. 148 | // 149 | // Wait for node to appear... 150 | select { 151 | case <-ctx.Done(): 152 | return nil, ctx.Err() 153 | case node := <-ch: 154 | return node, nil 155 | } 156 | } 157 | 158 | // backgroundNodesUpdate periodically checks list of registered nodes 159 | func (cl *Cluster[T]) backgroundNodesUpdate(ctx context.Context) { 160 | // initial update 161 | cl.updateNodes(ctx) 162 | 163 | ticker := time.NewTicker(cl.updateInterval) 164 | defer ticker.Stop() 165 | 166 | for { 167 | select { 168 | case <-ctx.Done(): 169 | return 170 | case <-ticker.C: 171 | cl.updateNodes(ctx) 172 | } 173 | } 174 | } 175 | 176 | // updateNodes performs a new round of cluster state check 177 | // and notifies all subscribers afterwards 178 | func (cl *Cluster[T]) updateNodes(ctx context.Context) { 179 | if cl.tracer.UpdateNodes != nil { 180 | cl.tracer.UpdateNodes() 181 | } 182 | 183 | ctx, cancel := context.WithTimeout(ctx, cl.updateTimeout) 184 | defer cancel() 185 | 186 | checked := checkNodes(ctx, cl.discoverer, cl.checker, cl.picker.CompareNodes, cl.tracer) 187 | cl.checkedNodes.Store(checked) 188 | 189 | if cl.tracer.NodesUpdated != nil { 190 | cl.tracer.NodesUpdated(checked) 191 | } 192 | 193 | cl.notifyUpdateSubscribers(checked) 194 | 195 | if cl.tracer.WaitersNotified != nil { 196 | cl.tracer.WaitersNotified() 197 | } 198 | } 199 | -------------------------------------------------------------------------------- /cluster_opts.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright 2020 YANDEX LLC 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | */ 16 | 17 | package hasql 18 | 19 | import "time" 20 | 21 | // ClusterOpt is a functional option type for Cluster constructor 22 | type ClusterOpt[T Querier] func(*Cluster[T]) 23 | 24 | // WithUpdateInterval sets interval between cluster state updates 25 | func WithUpdateInterval[T Querier](d time.Duration) ClusterOpt[T] { 26 | return func(cl *Cluster[T]) { 27 | cl.updateInterval = d 28 | } 29 | } 30 | 31 | // WithUpdateTimeout sets timeout for update of each node in cluster 32 | func WithUpdateTimeout[T Querier](d time.Duration) ClusterOpt[T] { 33 | return func(cl *Cluster[T]) { 34 | cl.updateTimeout = d 35 | } 36 | } 37 | 38 | // WithNodePicker sets algorithm for node selection (e.g. random, round robin etc.) 39 | func WithNodePicker[T Querier](picker NodePicker[T]) ClusterOpt[T] { 40 | return func(cl *Cluster[T]) { 41 | cl.picker = picker 42 | } 43 | } 44 | 45 | // WithTracer sets tracer for actions happening in the background 46 | func WithTracer[T Querier](tracer Tracer[T]) ClusterOpt[T] { 47 | return func(cl *Cluster[T]) { 48 | cl.tracer = tracer 49 | } 50 | } 51 | -------------------------------------------------------------------------------- /cluster_test.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright 2020 YANDEX LLC 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | */ 16 | 17 | package hasql 18 | 19 | import ( 20 | "database/sql" 21 | "io" 22 | "testing" 23 | 24 | "github.com/DATA-DOG/go-sqlmock" 25 | "github.com/stretchr/testify/assert" 26 | "github.com/stretchr/testify/require" 27 | ) 28 | 29 | func TestNewCluster(t *testing.T) { 30 | t.Run("no_nodes", func(t *testing.T) { 31 | cl, err := NewCluster[*sql.DB](nil, PostgreSQLChecker) 32 | assert.Nil(t, cl) 33 | assert.EqualError(t, err, "node discoverer required") 34 | }) 35 | 36 | t.Run("success", func(t *testing.T) { 37 | db, _, err := sqlmock.New() 38 | require.NoError(t, err) 39 | 40 | node := NewNode("shimba", db) 41 | 42 | cl, err := NewCluster(NewStaticNodeDiscoverer(node), PostgreSQLChecker) 43 | assert.NoError(t, err) 44 | assert.NotNil(t, cl) 45 | }) 46 | } 47 | 48 | func TestCluster_Close(t *testing.T) { 49 | t.Run("no_errors", func(t *testing.T) { 50 | db1, dbmock1, err := sqlmock.New() 51 | require.NoError(t, err) 52 | 53 | db2, dbmock2, err := sqlmock.New() 54 | require.NoError(t, err) 55 | 56 | // expect database client to be closed 57 | dbmock1.ExpectClose() 58 | dbmock2.ExpectClose() 59 | 60 | node1 := NewNode("shimba", db1) 61 | node2 := NewNode("boomba", db2) 62 | 63 | cl, err := NewCluster(NewStaticNodeDiscoverer(node1, node2), PostgreSQLChecker) 64 | require.NoError(t, err) 65 | 66 | cl.checkedNodes.Store(CheckedNodes[*sql.DB]{ 67 | discovered: []*Node[*sql.DB]{node1, node2}, 68 | }) 69 | 70 | assert.NoError(t, cl.Close()) 71 | }) 72 | 73 | t.Run("multiple_errors", func(t *testing.T) { 74 | db1, dbmock1, err := sqlmock.New() 75 | require.NoError(t, err) 76 | 77 | db2, dbmock2, err := sqlmock.New() 78 | require.NoError(t, err) 79 | 80 | // expect database client to be closed 81 | dbmock1.ExpectClose().WillReturnError(io.EOF) 82 | dbmock2.ExpectClose().WillReturnError(sql.ErrTxDone) 83 | 84 | node1 := NewNode("shimba", db1) 85 | node2 := NewNode("boomba", db2) 86 | 87 | cl, err := NewCluster(NewStaticNodeDiscoverer(node1, node2), PostgreSQLChecker) 88 | require.NoError(t, err) 89 | 90 | cl.checkedNodes.Store(CheckedNodes[*sql.DB]{ 91 | discovered: []*Node[*sql.DB]{node1, node2}, 92 | }) 93 | 94 | err = cl.Close() 95 | assert.ErrorIs(t, err, io.EOF) 96 | assert.ErrorIs(t, err, sql.ErrTxDone) 97 | }) 98 | } 99 | 100 | func TestCluster_Err(t *testing.T) { 101 | t.Run("no_error", func(t *testing.T) { 102 | cl := new(Cluster[*sql.DB]) 103 | cl.checkedNodes.Store(CheckedNodes[*sql.DB]{}) 104 | assert.NoError(t, cl.Err()) 105 | }) 106 | 107 | t.Run("has_error", func(t *testing.T) { 108 | checkedNodes := CheckedNodes[*sql.DB]{ 109 | err: NodeCheckErrors[*sql.DB]{ 110 | { 111 | node: &Node[*sql.DB]{ 112 | name: "shimba", 113 | db: new(sql.DB), 114 | }, 115 | err: io.EOF, 116 | }, 117 | }, 118 | } 119 | 120 | cl := new(Cluster[*sql.DB]) 121 | cl.checkedNodes.Store(checkedNodes) 122 | 123 | assert.ErrorIs(t, cl.Err(), io.EOF) 124 | }) 125 | } 126 | 127 | func TestCluster_Node(t *testing.T) { 128 | t.Run("no_nodes", func(t *testing.T) { 129 | cl := new(Cluster[*sql.DB]) 130 | cl.checkedNodes.Store(CheckedNodes[*sql.DB]{}) 131 | 132 | // all criteria must return nil node 133 | for i := Alive; i < maxNodeCriterion; i++ { 134 | node := cl.Node(i) 135 | assert.Nil(t, node) 136 | } 137 | }) 138 | 139 | t.Run("alive", func(t *testing.T) { 140 | node := &Node[*sql.DB]{ 141 | name: "shimba", 142 | db: new(sql.DB), 143 | } 144 | 145 | cl := new(Cluster[*sql.DB]) 146 | cl.picker = new(RandomNodePicker[*sql.DB]) 147 | cl.checkedNodes.Store(CheckedNodes[*sql.DB]{ 148 | alive: []CheckedNode[*sql.DB]{ 149 | { 150 | Node: node, 151 | Info: NodeInfo{ 152 | ClusterRole: NodeRolePrimary, 153 | }, 154 | }, 155 | }, 156 | }) 157 | 158 | assert.Equal(t, node, cl.Node(Alive)) 159 | }) 160 | 161 | t.Run("primary", func(t *testing.T) { 162 | node := &Node[*sql.DB]{ 163 | name: "shimba", 164 | db: new(sql.DB), 165 | } 166 | 167 | cl := new(Cluster[*sql.DB]) 168 | cl.picker = new(RandomNodePicker[*sql.DB]) 169 | cl.checkedNodes.Store(CheckedNodes[*sql.DB]{ 170 | primaries: []CheckedNode[*sql.DB]{ 171 | { 172 | Node: node, 173 | Info: NodeInfo{ 174 | ClusterRole: NodeRolePrimary, 175 | }, 176 | }, 177 | }, 178 | }) 179 | 180 | assert.Equal(t, node, cl.Node(Primary)) 181 | // we will return node on Prefer* criteria also 182 | assert.Equal(t, node, cl.Node(PreferPrimary)) 183 | assert.Equal(t, node, cl.Node(PreferStandby)) 184 | }) 185 | 186 | t.Run("standby", func(t *testing.T) { 187 | node := &Node[*sql.DB]{ 188 | name: "shimba", 189 | db: new(sql.DB), 190 | } 191 | 192 | cl := new(Cluster[*sql.DB]) 193 | cl.picker = new(RandomNodePicker[*sql.DB]) 194 | cl.checkedNodes.Store(CheckedNodes[*sql.DB]{ 195 | standbys: []CheckedNode[*sql.DB]{ 196 | { 197 | Node: node, 198 | Info: NodeInfo{ 199 | ClusterRole: NodeRoleStandby, 200 | }, 201 | }, 202 | }, 203 | }) 204 | 205 | assert.Equal(t, node, cl.Node(Standby)) 206 | // we will return node on Prefer* criteria also 207 | assert.Equal(t, node, cl.Node(PreferPrimary)) 208 | assert.Equal(t, node, cl.Node(PreferStandby)) 209 | }) 210 | 211 | t.Run("prefer_primary", func(t *testing.T) { 212 | node := &Node[*sql.DB]{ 213 | name: "shimba", 214 | db: new(sql.DB), 215 | } 216 | 217 | cl := new(Cluster[*sql.DB]) 218 | cl.picker = new(RandomNodePicker[*sql.DB]) 219 | 220 | // must pick from primaries 221 | cl.checkedNodes.Store(CheckedNodes[*sql.DB]{ 222 | primaries: []CheckedNode[*sql.DB]{ 223 | { 224 | Node: node, 225 | Info: NodeInfo{ 226 | ClusterRole: NodeRolePrimary, 227 | }, 228 | }, 229 | }, 230 | }) 231 | assert.Equal(t, node, cl.Node(PreferPrimary)) 232 | 233 | // must pick from standbys 234 | cl.checkedNodes.Store(CheckedNodes[*sql.DB]{ 235 | standbys: []CheckedNode[*sql.DB]{ 236 | { 237 | Node: node, 238 | Info: NodeInfo{ 239 | ClusterRole: NodeRoleStandby, 240 | }, 241 | }, 242 | }, 243 | }) 244 | assert.Equal(t, node, cl.Node(PreferPrimary)) 245 | }) 246 | 247 | t.Run("prefer_standby", func(t *testing.T) { 248 | node := &Node[*sql.DB]{ 249 | name: "shimba", 250 | db: new(sql.DB), 251 | } 252 | 253 | cl := new(Cluster[*sql.DB]) 254 | cl.picker = new(RandomNodePicker[*sql.DB]) 255 | 256 | // must pick from standbys 257 | cl.checkedNodes.Store(CheckedNodes[*sql.DB]{ 258 | standbys: []CheckedNode[*sql.DB]{ 259 | { 260 | Node: node, 261 | Info: NodeInfo{ 262 | ClusterRole: NodeRoleStandby, 263 | }, 264 | }, 265 | }, 266 | }) 267 | assert.Equal(t, node, cl.Node(PreferStandby)) 268 | 269 | // must pick from primaries 270 | cl.checkedNodes.Store(CheckedNodes[*sql.DB]{ 271 | primaries: []CheckedNode[*sql.DB]{ 272 | { 273 | Node: node, 274 | Info: NodeInfo{ 275 | ClusterRole: NodeRolePrimary, 276 | }, 277 | }, 278 | }, 279 | }) 280 | assert.Equal(t, node, cl.Node(PreferStandby)) 281 | }) 282 | } 283 | -------------------------------------------------------------------------------- /e2e_test.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright 2020 YANDEX LLC 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | */ 16 | 17 | package hasql_test 18 | 19 | import ( 20 | "context" 21 | "database/sql" 22 | "errors" 23 | "io" 24 | "slices" 25 | "sync/atomic" 26 | "testing" 27 | "time" 28 | 29 | "github.com/DATA-DOG/go-sqlmock" 30 | "github.com/stretchr/testify/assert" 31 | "github.com/stretchr/testify/require" 32 | 33 | "golang.yandex/hasql/v2" 34 | ) 35 | 36 | // TestEnd2End_AliveCluster setups 3 node cluster, waits for at least one 37 | // alive node, then picks primary and secondary node. Nodes are always alive. 38 | func TestEnd2End_AliveCluster(t *testing.T) { 39 | // create three database pools 40 | db1, mock1, err := sqlmock.New() 41 | require.NoError(t, err) 42 | db2, mock2, err := sqlmock.New() 43 | require.NoError(t, err) 44 | db3, mock3, err := sqlmock.New() 45 | require.NoError(t, err) 46 | 47 | // set db1 to be primary node 48 | mock1.ExpectQuery(`SELECT.*pg_is_in_recovery`). 49 | WillReturnRows(sqlmock. 50 | NewRows([]string{"role", "lag"}). 51 | AddRow(hasql.NodeRolePrimary, 0), 52 | ) 53 | 54 | // set db2 and db3 to be standby nodes 55 | mock2.ExpectQuery(`SELECT.*pg_is_in_recovery`). 56 | WillReturnRows(sqlmock. 57 | NewRows([]string{"role", "lag"}). 58 | AddRow(hasql.NodeRoleStandby, 0), 59 | ) 60 | mock3.ExpectQuery(`SELECT.*pg_is_in_recovery`). 61 | WillReturnRows(sqlmock. 62 | NewRows([]string{"role", "lag"}). 63 | AddRow(hasql.NodeRoleStandby, 10), 64 | ) 65 | 66 | // all pools must be closed in the end 67 | mock1.ExpectClose() 68 | mock2.ExpectClose() 69 | mock3.ExpectClose() 70 | 71 | // register pools as nodes 72 | node1 := hasql.NewNode("ololo", db1) 73 | node2 := hasql.NewNode("trololo", db2) 74 | node3 := hasql.NewNode("shimba", db3) 75 | discoverer := hasql.NewStaticNodeDiscoverer(node1, node2, node3) 76 | 77 | // create test cluster 78 | cl, err := hasql.NewCluster(discoverer, hasql.PostgreSQLChecker, 79 | hasql.WithUpdateInterval[*sql.DB](10*time.Millisecond), 80 | ) 81 | require.NoError(t, err) 82 | 83 | // close cluster and all underlying pools in the end 84 | defer func() { 85 | assert.NoError(t, cl.Close()) 86 | }() 87 | 88 | ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond) 89 | defer cancel() 90 | 91 | // wait for any alive node 92 | waitNode, err := cl.WaitForNode(ctx, hasql.Alive) 93 | assert.NoError(t, err) 94 | assert.Contains(t, []*hasql.Node[*sql.DB]{node1, node2, node3}, waitNode) 95 | 96 | // pick primary node 97 | primary := cl.Node(hasql.Primary) 98 | assert.Same(t, node1, primary) 99 | 100 | // pick standby node 101 | standby := cl.Node(hasql.Standby) 102 | assert.Contains(t, []*hasql.Node[*sql.DB]{node2, node3}, standby) 103 | } 104 | 105 | // TestEnd2End_SingleDeadNodeCluster setups 3 node cluster, waits for at least one 106 | // alive node, then picks primary and secondary node. One node is always dead. 107 | func TestEnd2End_SingleDeadNodeCluster(t *testing.T) { 108 | // create three database pools 109 | db1, mock1, err := sqlmock.New() 110 | require.NoError(t, err) 111 | db2, mock2, err := sqlmock.New() 112 | require.NoError(t, err) 113 | db3, mock3, err := sqlmock.New() 114 | require.NoError(t, err) 115 | 116 | // set db1 to be primary node 117 | mock1.ExpectQuery(`SELECT.*pg_is_in_recovery`). 118 | WillReturnRows(sqlmock. 119 | NewRows([]string{"role", "lag"}). 120 | AddRow(hasql.NodeRolePrimary, 0), 121 | ) 122 | // set db2 to be standby node 123 | mock2.ExpectQuery(`SELECT.*pg_is_in_recovery`). 124 | WillReturnRows(sqlmock. 125 | NewRows([]string{"role", "lag"}). 126 | AddRow(hasql.NodeRoleStandby, 0), 127 | ) 128 | // db3 will be always dead 129 | mock3.ExpectQuery(`SELECT.*pg_is_in_recovery`). 130 | WillDelayFor(time.Second). 131 | WillReturnError(io.EOF) 132 | 133 | // all pools must be closed in the end 134 | mock1.ExpectClose() 135 | mock2.ExpectClose() 136 | mock3.ExpectClose() 137 | 138 | // register pools as nodes 139 | node1 := hasql.NewNode("ololo", db1) 140 | node2 := hasql.NewNode("trololo", db2) 141 | node3 := hasql.NewNode("shimba", db3) 142 | discoverer := hasql.NewStaticNodeDiscoverer(node1, node2, node3) 143 | 144 | // create test cluster. 145 | cl, err := hasql.NewCluster(discoverer, hasql.PostgreSQLChecker, 146 | hasql.WithUpdateInterval[*sql.DB](10*time.Millisecond), 147 | hasql.WithUpdateTimeout[*sql.DB](50*time.Millisecond), 148 | // set node picker to round robin to guarantee iteration across all nodes 149 | hasql.WithNodePicker(new(hasql.RoundRobinNodePicker[*sql.DB])), 150 | ) 151 | require.NoError(t, err) 152 | 153 | // close cluster and all underlying pools in the end 154 | defer func() { 155 | assert.NoError(t, cl.Close()) 156 | }() 157 | 158 | // Set context timeout to be greater than cluster update interval and timeout. 159 | // If we set update timeout to be greater than wait context timeout 160 | // we will always receive context.DeadlineExceeded error as current cycle of update 161 | // will try to gather info about dead node (and thus update whole cluster state) 162 | // longer that we are waiting for node 163 | ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond) 164 | defer cancel() 165 | 166 | // wait for any alive node 167 | waitNode, err := cl.WaitForNode(ctx, hasql.Alive) 168 | assert.NoError(t, err) 169 | assert.Contains(t, []*hasql.Node[*sql.DB]{node1, node2}, waitNode) 170 | 171 | // pick primary node 172 | primary := cl.Node(hasql.Primary) 173 | assert.Same(t, node1, primary) 174 | 175 | // pick standby node multiple times to ensure 176 | // we always get alive standby node 177 | for range 100 { 178 | standby := cl.Node(hasql.Standby) 179 | assert.Same(t, node2, standby) 180 | } 181 | } 182 | 183 | // TestEnd2End_NoPrimaryCluster setups 3 node cluster, waits for at least one 184 | // alive node, then picks primary and secondary node. No alive primary nodes present. 185 | func TestEnd2End_NoPrimaryCluster(t *testing.T) { 186 | // create three database pools 187 | db1, mock1, err := sqlmock.New() 188 | require.NoError(t, err) 189 | db2, mock2, err := sqlmock.New() 190 | require.NoError(t, err) 191 | db3, mock3, err := sqlmock.New() 192 | require.NoError(t, err) 193 | 194 | // db1 is always dead 195 | mock1.ExpectQuery(`SELECT.*pg_is_in_recovery`). 196 | WillReturnError(io.EOF) 197 | // set db2 to be standby node 198 | mock2.ExpectQuery(`SELECT.*pg_is_in_recovery`). 199 | WillReturnRows(sqlmock. 200 | NewRows([]string{"role", "lag"}). 201 | AddRow(hasql.NodeRoleStandby, 10), 202 | ) 203 | // set db3 to be standby node 204 | mock3.ExpectQuery(`SELECT.*pg_is_in_recovery`). 205 | WillReturnRows(sqlmock. 206 | NewRows([]string{"role", "lag"}). 207 | AddRow(hasql.NodeRoleStandby, 0), 208 | ) 209 | 210 | // all pools must be closed in the end 211 | mock1.ExpectClose() 212 | mock2.ExpectClose() 213 | mock3.ExpectClose() 214 | 215 | // register pools as nodes 216 | node1 := hasql.NewNode("ololo", db1) 217 | node2 := hasql.NewNode("trololo", db2) 218 | node3 := hasql.NewNode("shimba", db3) 219 | discoverer := hasql.NewStaticNodeDiscoverer(node1, node2, node3) 220 | 221 | // create test cluster. 222 | cl, err := hasql.NewCluster(discoverer, hasql.PostgreSQLChecker, 223 | hasql.WithUpdateInterval[*sql.DB](10*time.Millisecond), 224 | ) 225 | require.NoError(t, err) 226 | 227 | // close cluster and all underlying pools in the end 228 | defer func() { 229 | assert.NoError(t, cl.Close()) 230 | }() 231 | 232 | ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond) 233 | defer cancel() 234 | 235 | // wait for any alive node 236 | waitNode, err := cl.WaitForNode(ctx, hasql.Alive) 237 | assert.NoError(t, err) 238 | assert.Contains(t, []*hasql.Node[*sql.DB]{node2, node3}, waitNode) 239 | 240 | // pick primary node 241 | primary := cl.Node(hasql.Primary) 242 | assert.Nil(t, primary) 243 | 244 | // pick standby node 245 | standby := cl.Node(hasql.Standby) 246 | assert.Contains(t, []*hasql.Node[*sql.DB]{node2, node3}, standby) 247 | 248 | // cluster must keep last check error 249 | assert.ErrorIs(t, cl.Err(), io.EOF) 250 | } 251 | 252 | // TestEnd2End_DeadCluster setups 3 node cluster. None node is alive. 253 | func TestEnd2End_DeadCluster(t *testing.T) { 254 | // create three database pools 255 | db1, mock1, err := sqlmock.New() 256 | require.NoError(t, err) 257 | db2, mock2, err := sqlmock.New() 258 | require.NoError(t, err) 259 | db3, mock3, err := sqlmock.New() 260 | require.NoError(t, err) 261 | 262 | // db1 is always dead 263 | mock1.ExpectQuery(`SELECT.*pg_is_in_recovery`). 264 | WillReturnError(io.EOF) 265 | // set db2 to be standby node 266 | mock2.ExpectQuery(`SELECT.*pg_is_in_recovery`). 267 | WillReturnError(io.EOF) 268 | // set db3 to be standby node 269 | mock3.ExpectQuery(`SELECT.*pg_is_in_recovery`). 270 | WillReturnError(io.EOF) 271 | 272 | // all pools must be closed in the end 273 | mock1.ExpectClose() 274 | mock2.ExpectClose() 275 | mock3.ExpectClose() 276 | 277 | // register pools as nodes 278 | node1 := hasql.NewNode("ololo", db1) 279 | node2 := hasql.NewNode("trololo", db2) 280 | node3 := hasql.NewNode("shimba", db3) 281 | discoverer := hasql.NewStaticNodeDiscoverer(node1, node2, node3) 282 | 283 | // create test cluster. 284 | cl, err := hasql.NewCluster(discoverer, hasql.PostgreSQLChecker, 285 | hasql.WithUpdateInterval[*sql.DB](10*time.Millisecond), 286 | hasql.WithUpdateTimeout[*sql.DB](50*time.Millisecond), 287 | ) 288 | require.NoError(t, err) 289 | 290 | // close cluster and all underlying pools in the end 291 | defer func() { 292 | assert.NoError(t, cl.Close()) 293 | }() 294 | 295 | // set context expiration to be greater than cluster refresh interval and timeout 296 | // to guarantee at least one cycle of state refresh 297 | ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond) 298 | defer cancel() 299 | 300 | // wait for any alive node 301 | waitNode, err := cl.WaitForNode(ctx, hasql.Alive) 302 | assert.ErrorIs(t, err, context.DeadlineExceeded) 303 | assert.Nil(t, waitNode) 304 | 305 | // pick primary node 306 | primary := cl.Node(hasql.Primary) 307 | assert.Nil(t, primary) 308 | 309 | // pick standby node 310 | standby := cl.Node(hasql.Standby) 311 | assert.Nil(t, standby) 312 | } 313 | 314 | // TestEnd2End_FlakyCluster setups 3 node cluster, waits for at least one 315 | // alive node, then picks primary and secondary node. 316 | // One node fails to report it's state between refresh intervals. 317 | func TestEnd2End_FlakyCluster(t *testing.T) { 318 | errIsPrimary := errors.New("primary node") 319 | errIsStandby := errors.New("standby node") 320 | 321 | sentinelErrChecker := func(ctx context.Context, q hasql.Querier) (hasql.NodeInfoProvider, error) { 322 | _, err := q.QueryContext(ctx, "report node pls") 323 | if errors.Is(err, errIsPrimary) { 324 | return hasql.NodeInfo{ClusterRole: hasql.NodeRolePrimary}, nil 325 | } 326 | if errors.Is(err, errIsStandby) { 327 | return hasql.NodeInfo{ClusterRole: hasql.NodeRoleStandby}, nil 328 | } 329 | return nil, err 330 | } 331 | 332 | // set db1 to be primary node 333 | // it will fail with error on every second attempt to query state 334 | var attempts uint32 335 | db1 := &mockQuerier{ 336 | queryFn: func(_ context.Context, _ string, _ ...any) (*sql.Rows, error) { 337 | call := atomic.AddUint32(&attempts, 1) 338 | if call%2 == 0 { 339 | return nil, io.EOF 340 | } 341 | return nil, errIsPrimary 342 | }, 343 | } 344 | 345 | // set db2 and db3 to be standbys 346 | db2 := &mockQuerier{ 347 | queryFn: func(_ context.Context, _ string, _ ...any) (*sql.Rows, error) { 348 | return nil, errIsStandby 349 | }, 350 | } 351 | db3 := &mockQuerier{ 352 | queryFn: func(_ context.Context, _ string, _ ...any) (*sql.Rows, error) { 353 | return nil, errIsStandby 354 | }, 355 | } 356 | 357 | // register pools as nodes 358 | node1 := hasql.NewNode("ololo", db1) 359 | node2 := hasql.NewNode("trololo", db2) 360 | node3 := hasql.NewNode("shimba", db3) 361 | discoverer := hasql.NewStaticNodeDiscoverer(node1, node2, node3) 362 | 363 | // create test cluster 364 | cl, err := hasql.NewCluster(discoverer, sentinelErrChecker, 365 | hasql.WithUpdateInterval[*mockQuerier](50*time.Millisecond), 366 | ) 367 | require.NoError(t, err) 368 | 369 | // close cluster and all underlying pools in the end 370 | defer func() { 371 | assert.NoError(t, cl.Close()) 372 | }() 373 | 374 | // wait for a long time 375 | ctx, cancel := context.WithTimeout(context.Background(), 500*time.Millisecond) 376 | defer cancel() 377 | 378 | // wait for any alive node 379 | waitNode, err := cl.WaitForNode(ctx, hasql.Alive) 380 | assert.NoError(t, err) 381 | assert.Contains(t, []*hasql.Node[*mockQuerier]{node1, node2, node3}, waitNode) 382 | 383 | // fetch nodes often 384 | ticker := time.NewTicker(10 * time.Millisecond) 385 | defer ticker.Stop() 386 | 387 | var primaryStates []string 388 | for { 389 | select { 390 | case <-ctx.Done(): 391 | // check that primary node has changed its state at least once 392 | expected := []string{"ololo", "", "ololo", "", "ololo", "", "ololo", "", "ololo", ""} 393 | assert.Equal(t, expected, slices.Compact(primaryStates)) 394 | // end test 395 | return 396 | case <-ticker.C: 397 | // pick primary node 398 | primary := cl.Node(hasql.Primary) 399 | // store current state for further checks 400 | var name string 401 | if primary != nil { 402 | name = primary.String() 403 | } 404 | primaryStates = append(primaryStates, name) 405 | 406 | // pick and check standby node 407 | standby := cl.Node(hasql.Standby) 408 | assert.Contains(t, []*hasql.Node[*mockQuerier]{node2, node3}, standby) 409 | } 410 | } 411 | 412 | } 413 | 414 | var _ hasql.Querier = (*mockQuerier)(nil) 415 | var _ io.Closer = (*mockQuerier)(nil) 416 | 417 | // mockQuerier returns fake SQL results to tests 418 | type mockQuerier struct { 419 | queryFn func(ctx context.Context, query string, args ...any) (*sql.Rows, error) 420 | queryRowFn func(ctx context.Context, query string, args ...any) *sql.Row 421 | closeFn func() error 422 | } 423 | 424 | func (m *mockQuerier) QueryContext(ctx context.Context, query string, args ...any) (*sql.Rows, error) { 425 | if m.queryFn != nil { 426 | return m.queryFn(ctx, query, args...) 427 | } 428 | return nil, nil 429 | } 430 | 431 | func (m *mockQuerier) QueryRowContext(ctx context.Context, query string, args ...any) *sql.Row { 432 | if m.queryRowFn != nil { 433 | return m.queryRowFn(ctx, query, args...) 434 | } 435 | return nil 436 | } 437 | 438 | func (m *mockQuerier) Close() error { 439 | if m.closeFn != nil { 440 | return m.closeFn() 441 | } 442 | return nil 443 | } 444 | -------------------------------------------------------------------------------- /error.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright 2020 YANDEX LLC 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | */ 16 | 17 | package hasql 18 | 19 | import ( 20 | "strings" 21 | ) 22 | 23 | // NodeCheckErrors is a set of checked nodes errors. 24 | // This type can be used in errors.As/Is as it implements errors.Unwrap method 25 | type NodeCheckErrors[T Querier] []NodeCheckError[T] 26 | 27 | func (n NodeCheckErrors[T]) Error() string { 28 | var b strings.Builder 29 | for i, err := range n { 30 | if i > 0 { 31 | b.WriteByte('\n') 32 | } 33 | b.WriteString(err.Error()) 34 | } 35 | return b.String() 36 | } 37 | 38 | // Unwrap is a helper for errors.Is/errors.As functions 39 | func (n NodeCheckErrors[T]) Unwrap() []error { 40 | errs := make([]error, len(n)) 41 | for i, err := range n { 42 | errs[i] = err 43 | } 44 | return errs 45 | } 46 | 47 | // NodeCheckError implements `error` and contains information about unsuccessful node check 48 | type NodeCheckError[T Querier] struct { 49 | node *Node[T] 50 | err error 51 | } 52 | 53 | // Node returns dead node instance 54 | func (n NodeCheckError[T]) Node() *Node[T] { 55 | return n.node 56 | } 57 | 58 | // Error implements `error` interface 59 | func (n NodeCheckError[T]) Error() string { 60 | if n.err == nil { 61 | return "" 62 | } 63 | return n.err.Error() 64 | } 65 | 66 | // Unwrap returns underlying error 67 | func (n NodeCheckError[T]) Unwrap() error { 68 | return n.err 69 | } 70 | -------------------------------------------------------------------------------- /example_test.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright 2020 YANDEX LLC 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | */ 16 | 17 | package hasql_test 18 | 19 | import ( 20 | "context" 21 | "database/sql" 22 | "time" 23 | 24 | "golang.yandex/hasql/v2" 25 | ) 26 | 27 | // ExampleCluster shows how to setup basic hasql cluster with some custom settings 28 | func ExampleCluster() { 29 | // open connections to database instances 30 | db1, err := sql.Open("pgx", "host1.example.com") 31 | if err != nil { 32 | panic(err) 33 | } 34 | db2, err := sql.Open("pgx", "host2.example.com") 35 | if err != nil { 36 | panic(err) 37 | } 38 | 39 | // register connections as nodes with some additional information 40 | nodes := []*hasql.Node[*sql.DB]{ 41 | hasql.NewNode("bear", db1), 42 | hasql.NewNode("battlestar galactica", db2), 43 | } 44 | 45 | // create NodeDiscoverer instance 46 | // here we use built-in StaticNodeDiscoverer which always returns all registered nodes 47 | discoverer := hasql.NewStaticNodeDiscoverer(nodes...) 48 | // use checker suitable for your database 49 | checker := hasql.PostgreSQLChecker 50 | // change default RandomNodePicker to RoundRobinNodePicker here 51 | picker := new(hasql.RoundRobinNodePicker[*sql.DB]) 52 | 53 | // create cluster instance using previously created discoverer, checker 54 | // and some additional options 55 | cl, err := hasql.NewCluster(discoverer, checker, 56 | // set custom picker via funcopt 57 | hasql.WithNodePicker(picker), 58 | // change interval of cluster state check 59 | hasql.WithUpdateInterval[*sql.DB](500*time.Millisecond), 60 | // change cluster check timeout value 61 | hasql.WithUpdateTimeout[*sql.DB](time.Second), 62 | ) 63 | if err != nil { 64 | panic(err) 65 | } 66 | 67 | // create context with timeout to wait for at least one alive node in cluster 68 | // note that context timeout value must be greater than cluster update interval + update timeout 69 | // otherwise you will always receive `context.DeadlineExceeded` error if one of cluster node is dead on startup 70 | ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) 71 | defer cancel() 72 | 73 | // wait for any alive node to guarantee that your application 74 | // does not start without database dependency 75 | // this step usually performed on application startup 76 | node, err := cl.WaitForNode(ctx, hasql.Alive) 77 | if err != nil { 78 | panic(err) 79 | } 80 | 81 | // pick standby node to perform query 82 | // always check node for nilness to avoid nil pointer dereference error 83 | // node object can be nil if no alive nodes for given criterion has been found 84 | node = cl.Node(hasql.Standby) 85 | if node == nil { 86 | panic("no alive standby available") 87 | } 88 | 89 | // get connection from node and perform desired action 90 | if err := node.DB().PingContext(ctx); err != nil { 91 | panic(err) 92 | } 93 | } 94 | -------------------------------------------------------------------------------- /go.mod: -------------------------------------------------------------------------------- 1 | module golang.yandex/hasql/v2 2 | 3 | go 1.22 4 | 5 | require ( 6 | github.com/DATA-DOG/go-sqlmock v1.5.2 7 | github.com/stretchr/testify v1.7.2 8 | ) 9 | 10 | require ( 11 | github.com/davecgh/go-spew v1.1.0 // indirect 12 | github.com/pmezard/go-difflib v1.0.0 // indirect 13 | gopkg.in/yaml.v3 v3.0.1 // indirect 14 | ) 15 | -------------------------------------------------------------------------------- /go.sum: -------------------------------------------------------------------------------- 1 | github.com/DATA-DOG/go-sqlmock v1.5.2 h1:OcvFkGmslmlZibjAjaHm3L//6LiuBgolP7OputlJIzU= 2 | github.com/DATA-DOG/go-sqlmock v1.5.2/go.mod h1:88MAG/4G7SMwSE3CeA0ZKzrT5CiOU3OJ+JlNzwDqpNU= 3 | github.com/davecgh/go-spew v1.1.0 h1:ZDRjVQ15GmhC3fiQ8ni8+OwkZQO4DARzQgrnXU1Liz8= 4 | github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= 5 | github.com/kisielk/sqlstruct v0.0.0-20201105191214-5f3e10d3ab46/go.mod h1:yyMNCyc/Ib3bDTKd379tNMpB/7/H5TjM2Y9QJ5THLbE= 6 | github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= 7 | github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= 8 | github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= 9 | github.com/stretchr/testify v1.7.2 h1:4jaiDzPyXQvSd7D0EjG45355tLlV3VOECpq10pLC+8s= 10 | github.com/stretchr/testify v1.7.2/go.mod h1:R6va5+xMeoiuVRoj+gSkQ7d3FALtqAAGI1FQKckRals= 11 | gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM= 12 | gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= 13 | gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= 14 | gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= 15 | -------------------------------------------------------------------------------- /mocks_test.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright 2020 YANDEX LLC 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | */ 16 | 17 | package hasql 18 | 19 | import ( 20 | "context" 21 | "database/sql" 22 | "io" 23 | "slices" 24 | ) 25 | 26 | var _ NodeDiscoverer[*mockQuerier] = (*mockNodeDiscoverer[*mockQuerier])(nil) 27 | 28 | // mockNodeDiscoverer returns stored results to tests 29 | type mockNodeDiscoverer[T Querier] struct { 30 | nodes []*Node[T] 31 | err error 32 | } 33 | 34 | func (e mockNodeDiscoverer[T]) DiscoverNodes(_ context.Context) ([]*Node[T], error) { 35 | return slices.Clone(e.nodes), e.err 36 | } 37 | 38 | var _ Querier = (*mockQuerier)(nil) 39 | var _ io.Closer = (*mockQuerier)(nil) 40 | 41 | // mockQuerier returns fake SQL results to tests 42 | type mockQuerier struct { 43 | name string 44 | queryFn func(ctx context.Context, query string, args ...any) (*sql.Rows, error) 45 | queryRowFn func(ctx context.Context, query string, args ...any) *sql.Row 46 | closeFn func() error 47 | } 48 | 49 | func (m *mockQuerier) QueryContext(ctx context.Context, query string, args ...any) (*sql.Rows, error) { 50 | if m.queryFn != nil { 51 | return m.queryFn(ctx, query, args...) 52 | } 53 | return nil, nil 54 | } 55 | 56 | func (m *mockQuerier) QueryRowContext(ctx context.Context, query string, args ...any) *sql.Row { 57 | if m.queryRowFn != nil { 58 | return m.queryRowFn(ctx, query, args...) 59 | } 60 | return nil 61 | } 62 | 63 | func (m *mockQuerier) Close() error { 64 | if m.closeFn != nil { 65 | return m.closeFn() 66 | } 67 | return nil 68 | } 69 | -------------------------------------------------------------------------------- /node.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright 2020 YANDEX LLC 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | */ 16 | 17 | package hasql 18 | 19 | // NodeStateCriterion represents a node selection criterion 20 | type NodeStateCriterion uint8 21 | 22 | const ( 23 | // Alive is a criterion to choose any alive node 24 | Alive NodeStateCriterion = iota + 1 25 | // Primary is a criterion to choose primary node 26 | Primary 27 | // Standby is a criterion to choose standby node 28 | Standby 29 | // PreferPrimary is a criterion to choose primary or any alive node 30 | PreferPrimary 31 | // PreferStandby is a criterion to choose standby or any alive node 32 | PreferStandby 33 | 34 | // maxNodeCriterion is for testing purposes only 35 | // all new criteria must be added above this constant 36 | maxNodeCriterion 37 | ) 38 | 39 | // Node holds reference to database connection pool with some additional data 40 | type Node[T Querier] struct { 41 | name string 42 | db T 43 | } 44 | 45 | // NewNode constructs node with given SQL querier 46 | func NewNode[T Querier](name string, db T) *Node[T] { 47 | return &Node[T]{name: name, db: db} 48 | } 49 | 50 | // DB returns node's database connection 51 | func (n *Node[T]) DB() T { 52 | return n.db 53 | } 54 | 55 | // String implements Stringer. 56 | // It uses name provided at construction to uniquely identify a single node 57 | func (n *Node[T]) String() string { 58 | return n.name 59 | } 60 | -------------------------------------------------------------------------------- /node_checker.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright 2020 YANDEX LLC 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | */ 16 | 17 | package hasql 18 | 19 | import ( 20 | "context" 21 | "math" 22 | "time" 23 | ) 24 | 25 | // NodeRole represents role of node in SQL cluster (usually primary/standby) 26 | type NodeRole uint8 27 | 28 | const ( 29 | // NodeRoleUnknown used to report node with unconventional role in cluster 30 | NodeRoleUnknown NodeRole = iota 31 | // NodeRolePrimary used to report node with primary role in cluster 32 | NodeRolePrimary 33 | // NodeRoleStandby used to report node with standby role in cluster 34 | NodeRoleStandby 35 | ) 36 | 37 | // NodeInfoProvider information about single cluster node 38 | type NodeInfoProvider interface { 39 | // Role reports role of node in cluster. 40 | // For SQL servers it is usually either primary or standby 41 | Role() NodeRole 42 | } 43 | 44 | // NodeInfo implements NodeInfoProvider with additional useful information 45 | var _ NodeInfoProvider = NodeInfo{} 46 | 47 | // NodeInfo contains various information about single cluster node 48 | type NodeInfo struct { 49 | // Role contains determined node's role in cluster 50 | ClusterRole NodeRole 51 | // Latency stores time that has been spent to send check request 52 | // and receive response from server 53 | NetworkLatency time.Duration 54 | // ReplicaLag represents how far behind is data on standby 55 | // in comparison to primary. As determination of real replication 56 | // lag is a tricky task and value type vary from one DBMS to another 57 | // (e.g. bytes count lag, time delta lag etc.) this field contains 58 | // abstract value for sorting purposes only 59 | ReplicaLag int 60 | } 61 | 62 | // Role reports determined role of node in cluster 63 | func (n NodeInfo) Role() NodeRole { 64 | return n.ClusterRole 65 | } 66 | 67 | // Latency reports time spend on query execution from client's point of view. 68 | // It can be used in LatencyNodePicker to determine node with fastest response time 69 | func (n NodeInfo) Latency() time.Duration { 70 | return n.NetworkLatency 71 | } 72 | 73 | // ReplicationLag reports data replication delta on standby. 74 | // It can be used in ReplicationNodePicker to determine node with most up-to-date data 75 | func (n NodeInfo) ReplicationLag() int { 76 | return n.ReplicaLag 77 | } 78 | 79 | // NodeChecker is a function that can perform request to SQL node and retrieve various information 80 | type NodeChecker func(context.Context, Querier) (NodeInfoProvider, error) 81 | 82 | // PostgreSQLChecker checks state on PostgreSQL node. 83 | // It reports appropriate information for PostgreSQL nodes version 10 and higher 84 | func PostgreSQLChecker(ctx context.Context, db Querier) (NodeInfoProvider, error) { 85 | start := time.Now() 86 | 87 | var role NodeRole 88 | var replicationLag *int 89 | err := db. 90 | QueryRowContext(ctx, ` 91 | SELECT 92 | ((pg_is_in_recovery())::int + 1) AS role, 93 | pg_last_wal_receive_lsn() - pg_last_wal_replay_lsn() AS replication_lag 94 | ; 95 | `). 96 | Scan(&role, &replicationLag) 97 | if err != nil { 98 | return nil, err 99 | } 100 | 101 | latency := time.Since(start) 102 | 103 | // determine proper replication lag value 104 | // by default we assume that replication is not started - hence maximum int value 105 | // see: https://www.postgresql.org/docs/current/functions-admin.html#FUNCTIONS-RECOVERY-CONTROL 106 | lag := math.MaxInt 107 | if replicationLag != nil { 108 | // use reported non-null replication lag 109 | lag = *replicationLag 110 | } 111 | if role == NodeRolePrimary { 112 | // primary node has no replication lag 113 | lag = 0 114 | } 115 | 116 | return NodeInfo{ 117 | ClusterRole: role, 118 | NetworkLatency: latency, 119 | ReplicaLag: lag, 120 | }, nil 121 | } 122 | 123 | // MySQLChecker checks state of MySQL node. 124 | // ATTENTION: database user must have REPLICATION CLIENT privilege to perform underlying query. 125 | func MySQLChecker(ctx context.Context, db Querier) (NodeInfoProvider, error) { 126 | start := time.Now() 127 | 128 | rows, err := db.QueryContext(ctx, "SHOW SLAVE STATUS") 129 | if err != nil { 130 | return nil, err 131 | } 132 | defer func() { _ = rows.Close() }() 133 | 134 | latency := time.Since(start) 135 | 136 | // only standby MySQL server will return rows for `SHOW SLAVE STATUS` query 137 | isStandby := rows.Next() 138 | // TODO: check SECONDS_BEHIND_MASTER row for "replication lag" 139 | 140 | if err := rows.Err(); err != nil { 141 | return nil, err 142 | } 143 | 144 | role := NodeRoleStandby 145 | lag := math.MaxInt 146 | if !isStandby { 147 | role = NodeRolePrimary 148 | lag = 0 149 | } 150 | 151 | return NodeInfo{ 152 | ClusterRole: role, 153 | NetworkLatency: latency, 154 | ReplicaLag: lag, 155 | }, nil 156 | } 157 | 158 | // MSSQLChecker checks state of MSSQL node 159 | func MSSQLChecker(ctx context.Context, db Querier) (NodeInfoProvider, error) { 160 | start := time.Now() 161 | 162 | var isPrimary bool 163 | err := db. 164 | QueryRowContext(ctx, ` 165 | SELECT 166 | IIF(count(database_guid) = 0, 'TRUE', 'FALSE') AS STATUS 167 | FROM sys.database_recovery_status 168 | WHERE database_guid IS NULL 169 | `). 170 | Scan(&isPrimary) 171 | if err != nil { 172 | return nil, err 173 | } 174 | 175 | latency := time.Since(start) 176 | role := NodeRoleStandby 177 | // TODO: proper replication lag calculation 178 | lag := math.MaxInt 179 | if isPrimary { 180 | role = NodeRolePrimary 181 | lag = 0 182 | } 183 | 184 | return NodeInfo{ 185 | ClusterRole: role, 186 | NetworkLatency: latency, 187 | ReplicaLag: lag, 188 | }, nil 189 | } 190 | -------------------------------------------------------------------------------- /node_discoverer.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright 2020 YANDEX LLC 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | */ 16 | 17 | package hasql 18 | 19 | import ( 20 | "context" 21 | "database/sql" 22 | ) 23 | 24 | // NodeDiscoverer represents a provider of cluster nodes list. 25 | // NodeDiscoverer must node check nodes liveness or role, just return all nodes registered in cluster 26 | type NodeDiscoverer[T Querier] interface { 27 | // DiscoverNodes returns list of nodes registered in cluster 28 | DiscoverNodes(context.Context) ([]*Node[T], error) 29 | } 30 | 31 | // StaticNodeDiscoverer implements NodeDiscoverer 32 | var _ NodeDiscoverer[*sql.DB] = (*StaticNodeDiscoverer[*sql.DB])(nil) 33 | 34 | // StaticNodeDiscoverer always returns list of provided nodes 35 | type StaticNodeDiscoverer[T Querier] struct { 36 | nodes []*Node[T] 37 | } 38 | 39 | // NewStaticNodeDiscoverer returns new staticNodeDiscoverer instance 40 | func NewStaticNodeDiscoverer[T Querier](nodes ...*Node[T]) StaticNodeDiscoverer[T] { 41 | return StaticNodeDiscoverer[T]{nodes: nodes} 42 | } 43 | 44 | // DiscoverNodes returns static list of nodes from StaticNodeDiscoverer 45 | func (s StaticNodeDiscoverer[T]) DiscoverNodes(_ context.Context) ([]*Node[T], error) { 46 | return s.nodes, nil 47 | } 48 | -------------------------------------------------------------------------------- /node_discoverer_test.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright 2020 YANDEX LLC 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | */ 16 | 17 | package hasql 18 | 19 | import ( 20 | "context" 21 | "database/sql" 22 | "testing" 23 | 24 | "github.com/stretchr/testify/assert" 25 | ) 26 | 27 | func TestNewStaticNodeDiscoverer(t *testing.T) { 28 | node1 := NewNode("shimba", new(sql.DB)) 29 | node2 := NewNode("boomba", new(sql.DB)) 30 | 31 | d := NewStaticNodeDiscoverer(node1, node2) 32 | expected := StaticNodeDiscoverer[*sql.DB]{ 33 | nodes: []*Node[*sql.DB]{node1, node2}, 34 | } 35 | 36 | assert.Equal(t, expected, d) 37 | } 38 | 39 | func TestStaticNodeDiscoverer_DiscoverNodes(t *testing.T) { 40 | node1 := NewNode("shimba", new(sql.DB)) 41 | node2 := NewNode("boomba", new(sql.DB)) 42 | 43 | d := NewStaticNodeDiscoverer(node1, node2) 44 | 45 | discovered, err := d.DiscoverNodes(context.Background()) 46 | assert.NoError(t, err) 47 | assert.Equal(t, []*Node[*sql.DB]{node1, node2}, discovered) 48 | } 49 | -------------------------------------------------------------------------------- /node_picker.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright 2020 YANDEX LLC 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | */ 16 | 17 | package hasql 18 | 19 | import ( 20 | "database/sql" 21 | "math/rand/v2" 22 | "sync/atomic" 23 | "time" 24 | ) 25 | 26 | // NodePicker decides which node must be used from given set. 27 | // It also provides a comparer to be used to pre-sort nodes for better performance 28 | type NodePicker[T Querier] interface { 29 | // PickNode returns a single node from given set 30 | PickNode(nodes []CheckedNode[T]) CheckedNode[T] 31 | // CompareNodes is a comparison function to be used to sort checked nodes 32 | CompareNodes(a, b CheckedNode[T]) int 33 | } 34 | 35 | // RandomNodePicker implements NodePicker 36 | var _ NodePicker[*sql.DB] = (*RandomNodePicker[*sql.DB])(nil) 37 | 38 | // RandomNodePicker returns random node on each call and does not sort checked nodes 39 | type RandomNodePicker[T Querier] struct{} 40 | 41 | // PickNode returns random node from picker 42 | func (*RandomNodePicker[T]) PickNode(nodes []CheckedNode[T]) CheckedNode[T] { 43 | return nodes[rand.IntN(len(nodes))] 44 | } 45 | 46 | // CompareNodes always treats nodes as equal, effectively not changing nodes order 47 | func (*RandomNodePicker[T]) CompareNodes(_, _ CheckedNode[T]) int { 48 | return 0 49 | } 50 | 51 | // RoundRobinNodePicker implements NodePicker 52 | var _ NodePicker[*sql.DB] = (*RoundRobinNodePicker[*sql.DB])(nil) 53 | 54 | // RoundRobinNodePicker returns next node based on Round Robin algorithm and tries to preserve nodes order across checks 55 | type RoundRobinNodePicker[T Querier] struct { 56 | idx uint32 57 | } 58 | 59 | // PickNode returns next node in Round-Robin sequence 60 | func (r *RoundRobinNodePicker[T]) PickNode(nodes []CheckedNode[T]) CheckedNode[T] { 61 | n := atomic.AddUint32(&r.idx, 1) 62 | return nodes[(int(n)-1)%len(nodes)] 63 | } 64 | 65 | // CompareNodes performs lexicographical comparison of two nodes 66 | func (r *RoundRobinNodePicker[T]) CompareNodes(a, b CheckedNode[T]) int { 67 | aName, bName := a.Node.String(), b.Node.String() 68 | if aName < bName { 69 | return -1 70 | } 71 | if aName > bName { 72 | return 1 73 | } 74 | return 0 75 | } 76 | 77 | // LatencyNodePicker implements NodePicker 78 | var _ NodePicker[*sql.DB] = (*LatencyNodePicker[*sql.DB])(nil) 79 | 80 | // LatencyNodePicker returns node with least latency and sorts checked nodes by reported latency ascending. 81 | // WARNING: This picker requires that NodeInfoProvider can report node's network latency otherwise code will panic! 82 | type LatencyNodePicker[T Querier] struct{} 83 | 84 | // PickNode returns node with least network latency 85 | func (*LatencyNodePicker[T]) PickNode(nodes []CheckedNode[T]) CheckedNode[T] { 86 | return nodes[0] 87 | } 88 | 89 | // CompareNodes performs nodes comparison based on reported network latency 90 | func (*LatencyNodePicker[T]) CompareNodes(a, b CheckedNode[T]) int { 91 | aLatency := a.Info.(interface{ Latency() time.Duration }).Latency() 92 | bLatency := b.Info.(interface{ Latency() time.Duration }).Latency() 93 | 94 | if aLatency < bLatency { 95 | return -1 96 | } 97 | if aLatency > bLatency { 98 | return 1 99 | } 100 | return 0 101 | } 102 | 103 | // ReplicationNodePicker implements NodePicker 104 | var _ NodePicker[*sql.DB] = (*ReplicationNodePicker[*sql.DB])(nil) 105 | 106 | // ReplicationNodePicker returns node with smallest replication lag and sorts checked nodes by reported replication lag ascending. 107 | // Note that replication lag reported by checkers can vastly differ from the real situation on standby server. 108 | // WARNING: This picker requires that NodeInfoProvider can report node's replication lag otherwise code will panic! 109 | type ReplicationNodePicker[T Querier] struct{} 110 | 111 | // PickNode returns node with lowest replication lag value 112 | func (*ReplicationNodePicker[T]) PickNode(nodes []CheckedNode[T]) CheckedNode[T] { 113 | return nodes[0] 114 | } 115 | 116 | // CompareNodes performs nodes comparison based on reported replication lag 117 | func (*ReplicationNodePicker[T]) CompareNodes(a, b CheckedNode[T]) int { 118 | aLag := a.Info.(interface{ ReplicationLag() int }).ReplicationLag() 119 | bLag := b.Info.(interface{ ReplicationLag() int }).ReplicationLag() 120 | 121 | if aLag < bLag { 122 | return -1 123 | } 124 | if aLag > bLag { 125 | return 1 126 | } 127 | return 0 128 | } 129 | -------------------------------------------------------------------------------- /node_picker_test.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright 2020 YANDEX LLC 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | */ 16 | 17 | package hasql 18 | 19 | import ( 20 | "database/sql" 21 | "testing" 22 | "time" 23 | 24 | "github.com/stretchr/testify/assert" 25 | ) 26 | 27 | func TestRandomNodePicker(t *testing.T) { 28 | t.Run("pick_node", func(t *testing.T) { 29 | nodes := []CheckedNode[*sql.DB]{ 30 | { 31 | Node: NewNode("shimba", (*sql.DB)(nil)), 32 | }, 33 | { 34 | Node: NewNode("boomba", (*sql.DB)(nil)), 35 | }, 36 | { 37 | Node: NewNode("looken", (*sql.DB)(nil)), 38 | }, 39 | } 40 | 41 | np := new(RandomNodePicker[*sql.DB]) 42 | 43 | pickedNodes := make(map[string]struct{}) 44 | for range 100 { 45 | pickedNodes[np.PickNode(nodes).Node.String()] = struct{}{} 46 | } 47 | expectedNodes := map[string]struct{}{"boomba": {}, "looken": {}, "shimba": {}} 48 | 49 | assert.Equal(t, expectedNodes, pickedNodes) 50 | }) 51 | 52 | t.Run("compare_nodes", func(t *testing.T) { 53 | a := CheckedNode[*sql.DB]{ 54 | Node: NewNode("shimba", (*sql.DB)(nil)), 55 | Info: NodeInfo{ 56 | ClusterRole: NodeRolePrimary, 57 | NetworkLatency: 10 * time.Millisecond, 58 | ReplicaLag: 1, 59 | }, 60 | } 61 | 62 | b := CheckedNode[*sql.DB]{ 63 | Node: NewNode("boomba", (*sql.DB)(nil)), 64 | Info: NodeInfo{ 65 | ClusterRole: NodeRoleStandby, 66 | NetworkLatency: 20 * time.Millisecond, 67 | ReplicaLag: 2, 68 | }, 69 | } 70 | 71 | np := new(RandomNodePicker[*sql.DB]) 72 | 73 | for _, nodeA := range []CheckedNode[*sql.DB]{a, b} { 74 | for _, nodeB := range []CheckedNode[*sql.DB]{a, b} { 75 | assert.Equal(t, 0, np.CompareNodes(nodeA, nodeB)) 76 | } 77 | } 78 | }) 79 | } 80 | 81 | func TestRoundRobinNodePicker(t *testing.T) { 82 | t.Run("pick_node", func(t *testing.T) { 83 | nodes := []CheckedNode[*sql.DB]{ 84 | { 85 | Node: NewNode("shimba", (*sql.DB)(nil)), 86 | }, 87 | { 88 | Node: NewNode("boomba", (*sql.DB)(nil)), 89 | }, 90 | { 91 | Node: NewNode("looken", (*sql.DB)(nil)), 92 | }, 93 | { 94 | Node: NewNode("tooken", (*sql.DB)(nil)), 95 | }, 96 | { 97 | Node: NewNode("chicken", (*sql.DB)(nil)), 98 | }, 99 | { 100 | Node: NewNode("cooken", (*sql.DB)(nil)), 101 | }, 102 | } 103 | 104 | np := new(RoundRobinNodePicker[*sql.DB]) 105 | 106 | var pickedNodes []string 107 | for range len(nodes) * 3 { 108 | pickedNodes = append(pickedNodes, np.PickNode(nodes).Node.String()) 109 | } 110 | 111 | expectedNodes := []string{ 112 | "shimba", "boomba", "looken", "tooken", "chicken", "cooken", 113 | "shimba", "boomba", "looken", "tooken", "chicken", "cooken", 114 | "shimba", "boomba", "looken", "tooken", "chicken", "cooken", 115 | } 116 | assert.Equal(t, expectedNodes, pickedNodes) 117 | }) 118 | 119 | t.Run("compare_nodes", func(t *testing.T) { 120 | a := CheckedNode[*sql.DB]{ 121 | Node: NewNode("shimba", (*sql.DB)(nil)), 122 | Info: NodeInfo{ 123 | ClusterRole: NodeRolePrimary, 124 | NetworkLatency: 10 * time.Millisecond, 125 | ReplicaLag: 1, 126 | }, 127 | } 128 | 129 | b := CheckedNode[*sql.DB]{ 130 | Node: NewNode("boomba", (*sql.DB)(nil)), 131 | Info: NodeInfo{ 132 | ClusterRole: NodeRoleStandby, 133 | NetworkLatency: 20 * time.Millisecond, 134 | ReplicaLag: 2, 135 | }, 136 | } 137 | 138 | np := new(RoundRobinNodePicker[*sql.DB]) 139 | 140 | assert.Equal(t, 1, np.CompareNodes(a, b)) 141 | assert.Equal(t, -1, np.CompareNodes(b, a)) 142 | assert.Equal(t, 0, np.CompareNodes(a, a)) 143 | assert.Equal(t, 0, np.CompareNodes(b, b)) 144 | }) 145 | } 146 | 147 | func TestLatencyNodePicker(t *testing.T) { 148 | t.Run("pick_node", func(t *testing.T) { 149 | nodes := []CheckedNode[*sql.DB]{ 150 | { 151 | Node: NewNode("shimba", (*sql.DB)(nil)), 152 | }, 153 | { 154 | Node: NewNode("boomba", (*sql.DB)(nil)), 155 | }, 156 | { 157 | Node: NewNode("looken", (*sql.DB)(nil)), 158 | }, 159 | { 160 | Node: NewNode("tooken", (*sql.DB)(nil)), 161 | }, 162 | { 163 | Node: NewNode("chicken", (*sql.DB)(nil)), 164 | }, 165 | { 166 | Node: NewNode("cooken", (*sql.DB)(nil)), 167 | }, 168 | } 169 | 170 | np := new(LatencyNodePicker[*sql.DB]) 171 | 172 | pickedNodes := make(map[string]struct{}) 173 | for range 100 { 174 | pickedNodes[np.PickNode(nodes).Node.String()] = struct{}{} 175 | } 176 | 177 | expectedNodes := map[string]struct{}{ 178 | "shimba": {}, 179 | } 180 | assert.Equal(t, expectedNodes, pickedNodes) 181 | }) 182 | 183 | t.Run("compare_nodes", func(t *testing.T) { 184 | a := CheckedNode[*sql.DB]{ 185 | Node: NewNode("shimba", (*sql.DB)(nil)), 186 | Info: NodeInfo{ 187 | ClusterRole: NodeRolePrimary, 188 | NetworkLatency: 10 * time.Millisecond, 189 | ReplicaLag: 1, 190 | }, 191 | } 192 | 193 | b := CheckedNode[*sql.DB]{ 194 | Node: NewNode("boomba", (*sql.DB)(nil)), 195 | Info: NodeInfo{ 196 | ClusterRole: NodeRoleStandby, 197 | NetworkLatency: 20 * time.Millisecond, 198 | ReplicaLag: 2, 199 | }, 200 | } 201 | 202 | np := new(LatencyNodePicker[*sql.DB]) 203 | 204 | assert.Equal(t, -1, np.CompareNodes(a, b)) 205 | assert.Equal(t, 1, np.CompareNodes(b, a)) 206 | assert.Equal(t, 0, np.CompareNodes(a, a)) 207 | assert.Equal(t, 0, np.CompareNodes(b, b)) 208 | }) 209 | } 210 | 211 | func TestReplicationNodePicker(t *testing.T) { 212 | t.Run("pick_node", func(t *testing.T) { 213 | nodes := []CheckedNode[*sql.DB]{ 214 | { 215 | Node: NewNode("shimba", (*sql.DB)(nil)), 216 | }, 217 | { 218 | Node: NewNode("boomba", (*sql.DB)(nil)), 219 | }, 220 | { 221 | Node: NewNode("looken", (*sql.DB)(nil)), 222 | }, 223 | { 224 | Node: NewNode("tooken", (*sql.DB)(nil)), 225 | }, 226 | { 227 | Node: NewNode("chicken", (*sql.DB)(nil)), 228 | }, 229 | { 230 | Node: NewNode("cooken", (*sql.DB)(nil)), 231 | }, 232 | } 233 | 234 | np := new(ReplicationNodePicker[*sql.DB]) 235 | 236 | pickedNodes := make(map[string]struct{}) 237 | for range 100 { 238 | pickedNodes[np.PickNode(nodes).Node.String()] = struct{}{} 239 | } 240 | 241 | expectedNodes := map[string]struct{}{ 242 | "shimba": {}, 243 | } 244 | assert.Equal(t, expectedNodes, pickedNodes) 245 | }) 246 | 247 | t.Run("compare_nodes", func(t *testing.T) { 248 | a := CheckedNode[*sql.DB]{ 249 | Node: NewNode("shimba", (*sql.DB)(nil)), 250 | Info: NodeInfo{ 251 | ClusterRole: NodeRolePrimary, 252 | NetworkLatency: 10 * time.Millisecond, 253 | ReplicaLag: 1, 254 | }, 255 | } 256 | 257 | b := CheckedNode[*sql.DB]{ 258 | Node: NewNode("boomba", (*sql.DB)(nil)), 259 | Info: NodeInfo{ 260 | ClusterRole: NodeRoleStandby, 261 | NetworkLatency: 20 * time.Millisecond, 262 | ReplicaLag: 2, 263 | }, 264 | } 265 | 266 | np := new(ReplicationNodePicker[*sql.DB]) 267 | 268 | assert.Equal(t, -1, np.CompareNodes(a, b)) 269 | assert.Equal(t, 1, np.CompareNodes(b, a)) 270 | assert.Equal(t, 0, np.CompareNodes(a, a)) 271 | assert.Equal(t, 0, np.CompareNodes(b, b)) 272 | }) 273 | } 274 | -------------------------------------------------------------------------------- /notify.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright 2020 YANDEX LLC 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | */ 16 | 17 | package hasql 18 | 19 | // updateSubscriber represents a waiter for newly checked node event 20 | type updateSubscriber[T Querier] struct { 21 | ch chan *Node[T] 22 | criterion NodeStateCriterion 23 | } 24 | 25 | // addUpdateSubscriber adds new subscriber to notification pool 26 | func (cl *Cluster[T]) addUpdateSubscriber(criterion NodeStateCriterion) <-chan *Node[T] { 27 | // buffered channel is essential 28 | // read WaitForNode function for more information 29 | ch := make(chan *Node[T], 1) 30 | cl.subscribersMu.Lock() 31 | defer cl.subscribersMu.Unlock() 32 | cl.subscribers = append(cl.subscribers, updateSubscriber[T]{ch: ch, criterion: criterion}) 33 | return ch 34 | } 35 | 36 | // notifyUpdateSubscribers sends appropriate nodes to registered subscribers. 37 | // This function uses newly checked nodes to avoid race conditions 38 | func (cl *Cluster[T]) notifyUpdateSubscribers(nodes CheckedNodes[T]) { 39 | cl.subscribersMu.Lock() 40 | defer cl.subscribersMu.Unlock() 41 | 42 | if len(cl.subscribers) == 0 { 43 | return 44 | } 45 | 46 | var nodelessWaiters []updateSubscriber[T] 47 | // Notify all waiters 48 | for _, subscriber := range cl.subscribers { 49 | node := pickNodeByCriterion(nodes, cl.picker, subscriber.criterion) 50 | if node == nil { 51 | // Put waiter back 52 | nodelessWaiters = append(nodelessWaiters, subscriber) 53 | continue 54 | } 55 | 56 | // We won't block here, read addUpdateWaiter function for more information 57 | subscriber.ch <- node 58 | // No need to close channel since we write only once and forget it so does the 'client' 59 | } 60 | 61 | cl.subscribers = nodelessWaiters 62 | } 63 | -------------------------------------------------------------------------------- /sql.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright 2020 YANDEX LLC 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | */ 16 | 17 | package hasql 18 | 19 | import ( 20 | "context" 21 | "database/sql" 22 | ) 23 | 24 | // Querier describes abstract base SQL client such as database/sql.DB. 25 | // Most of database/sql compatible third-party libraries already implement it 26 | type Querier interface { 27 | // QueryRowContext executes a query that is expected to return at most one row 28 | QueryRowContext(ctx context.Context, query string, args ...any) *sql.Row 29 | // QueryContext executes a query that returns rows, typically a SELECT 30 | QueryContext(ctx context.Context, query string, args ...any) (*sql.Rows, error) 31 | } 32 | -------------------------------------------------------------------------------- /trace.go: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright 2020 YANDEX LLC 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | */ 16 | 17 | package hasql 18 | 19 | // Tracer is a set of hooks to be called at various stages of background nodes status update. 20 | // Any particular hook may be nil. Functions may be called concurrently from different goroutines. 21 | type Tracer[T Querier] struct { 22 | // UpdateNodes is called when before updating nodes status. 23 | UpdateNodes func() 24 | // NodesUpdated is called after all nodes are updated. The nodes is a list of currently alive nodes. 25 | NodesUpdated func(nodes CheckedNodes[T]) 26 | // NodeDead is called when it is determined that specified node is dead. 27 | NodeDead func(err error) 28 | // NodeAlive is called when it is determined that specified node is alive. 29 | NodeAlive func(node CheckedNode[T]) 30 | // WaitersNotified is called when callers of 'WaitForNode' function have been notified. 31 | WaitersNotified func() 32 | } 33 | --------------------------------------------------------------------------------