├── .gitignore ├── BENCHMARK.md ├── CHANGELOG.md ├── LICENSE ├── README.md ├── main.go ├── onecache ├── config.go ├── memcache.go ├── memcache_handlers.go ├── memcache_test.go ├── onecache.go ├── onecache_rpc.go ├── onecache_rpc_test.go ├── onecache_test.go ├── pool.go ├── replicator.go ├── replicator_test.go ├── ring.go ├── ring_test.go ├── rpc_dispatch.go └── test_util.go └── ttlstore ├── ttl_heap.go ├── ttlstore.go ├── ttlstore_test.go └── util.go /.gitignore: -------------------------------------------------------------------------------- 1 | # Compiled Object files, Static and Dynamic libs (Shared Objects) 2 | *.o 3 | *.a 4 | *.so 5 | 6 | # Folders 7 | _obj 8 | _test 9 | 10 | # Architecture specific extensions/prefixes 11 | *.[568vq] 12 | [568vq].out 13 | 14 | *.cgo1.go 15 | *.cgo2.c 16 | _cgo_defun.c 17 | _cgo_gotypes.go 18 | _cgo_export.* 19 | 20 | _testmain.go 21 | 22 | *.exe 23 | *.test 24 | *.prof 25 | -------------------------------------------------------------------------------- /BENCHMARK.md: -------------------------------------------------------------------------------- 1 | Single-node benchmark done using this [tool](https://github.com/antirez/mc-benchmark/) on my laptop: 2 | 3 | # v0.1 4 | OneCache: 5 | ``` 6 | ./mc-benchmark -n 50000 7 | ====== SET ====== 8 | 50000 requests completed in 1.16 seconds 9 | 50 parallel clients 10 | 3 bytes payload 11 | keep alive: 1 12 | 13 | 4.49% <= 0 milliseconds 14 | 81.70% <= 1 milliseconds 15 | 98.84% <= 2 milliseconds 16 | 99.80% <= 3 milliseconds 17 | 99.97% <= 4 milliseconds 18 | 99.99% <= 5 milliseconds 19 | 99.99% <= 6 milliseconds 20 | 100.00% <= 7 milliseconds 21 | 42992.26 requests per second 22 | 23 | ====== GET ====== 24 | 50007 requests completed in 1.27 seconds 25 | 50 parallel clients 26 | 3 bytes payload 27 | keep alive: 1 28 | 29 | 3.90% <= 0 milliseconds 30 | 73.76% <= 1 milliseconds 31 | 97.94% <= 2 milliseconds 32 | 99.52% <= 3 milliseconds 33 | 99.93% <= 4 milliseconds 34 | 99.97% <= 5 milliseconds 35 | 99.98% <= 6 milliseconds 36 | 99.99% <= 7 milliseconds 37 | 99.99% <= 8 milliseconds 38 | 100.00% <= 9 milliseconds 39 | 39375.59 requests per second 40 | ``` 41 | 42 | memcached: 43 | ``` 44 | ./mc-benchmark -n 50000 45 | ====== SET ====== 46 | 50000 requests completed in 1.00 seconds 47 | 50 parallel clients 48 | 3 bytes payload 49 | keep alive: 1 50 | 51 | 5.33% <= 0 milliseconds 52 | 95.73% <= 1 milliseconds 53 | 99.58% <= 2 milliseconds 54 | 99.95% <= 3 milliseconds 55 | 99.99% <= 4 milliseconds 56 | 100.00% <= 5 milliseconds 57 | 49800.80 requests per second 58 | 59 | ====== GET ====== 60 | 50000 requests completed in 1.09 seconds 61 | 50 parallel clients 62 | 3 bytes payload 63 | keep alive: 1 64 | 65 | 7.93% <= 0 milliseconds 66 | 86.05% <= 1 milliseconds 67 | 98.86% <= 2 milliseconds 68 | 99.79% <= 3 milliseconds 69 | 99.90% <= 4 milliseconds 70 | 99.90% <= 5 milliseconds 71 | 99.92% <= 8 milliseconds 72 | 99.95% <= 9 milliseconds 73 | 99.97% <= 10 milliseconds 74 | 99.99% <= 11 milliseconds 75 | 100.00% <= 12 milliseconds 76 | 100.00% <= 13 milliseconds 77 | 45703.84 requests per second 78 | ``` 79 | -------------------------------------------------------------------------------- /CHANGELOG.md: -------------------------------------------------------------------------------- 1 | ## 0.1.0 (January 18, 2015) 2 | 3 | * Initial release 4 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Mozilla Public License Version 2.0 2 | ================================== 3 | 4 | 1. Definitions 5 | -------------- 6 | 7 | 1.1. "Contributor" 8 | means each individual or legal entity that creates, contributes to 9 | the creation of, or owns Covered Software. 10 | 11 | 1.2. "Contributor Version" 12 | means the combination of the Contributions of others (if any) used 13 | by a Contributor and that particular Contributor's Contribution. 14 | 15 | 1.3. "Contribution" 16 | means Covered Software of a particular Contributor. 17 | 18 | 1.4. "Covered Software" 19 | means Source Code Form to which the initial Contributor has attached 20 | the notice in Exhibit A, the Executable Form of such Source Code 21 | Form, and Modifications of such Source Code Form, in each case 22 | including portions thereof. 23 | 24 | 1.5. "Incompatible With Secondary Licenses" 25 | means 26 | 27 | (a) that the initial Contributor has attached the notice described 28 | in Exhibit B to the Covered Software; or 29 | 30 | (b) that the Covered Software was made available under the terms of 31 | version 1.1 or earlier of the License, but not also under the 32 | terms of a Secondary License. 33 | 34 | 1.6. "Executable Form" 35 | means any form of the work other than Source Code Form. 36 | 37 | 1.7. "Larger Work" 38 | means a work that combines Covered Software with other material, in 39 | a separate file or files, that is not Covered Software. 40 | 41 | 1.8. "License" 42 | means this document. 43 | 44 | 1.9. "Licensable" 45 | means having the right to grant, to the maximum extent possible, 46 | whether at the time of the initial grant or subsequently, any and 47 | all of the rights conveyed by this License. 48 | 49 | 1.10. "Modifications" 50 | means any of the following: 51 | 52 | (a) any file in Source Code Form that results from an addition to, 53 | deletion from, or modification of the contents of Covered 54 | Software; or 55 | 56 | (b) any new file in Source Code Form that contains any Covered 57 | Software. 58 | 59 | 1.11. "Patent Claims" of a Contributor 60 | means any patent claim(s), including without limitation, method, 61 | process, and apparatus claims, in any patent Licensable by such 62 | Contributor that would be infringed, but for the grant of the 63 | License, by the making, using, selling, offering for sale, having 64 | made, import, or transfer of either its Contributions or its 65 | Contributor Version. 66 | 67 | 1.12. "Secondary License" 68 | means either the GNU General Public License, Version 2.0, the GNU 69 | Lesser General Public License, Version 2.1, the GNU Affero General 70 | Public License, Version 3.0, or any later versions of those 71 | licenses. 72 | 73 | 1.13. "Source Code Form" 74 | means the form of the work preferred for making modifications. 75 | 76 | 1.14. "You" (or "Your") 77 | means an individual or a legal entity exercising rights under this 78 | License. For legal entities, "You" includes any entity that 79 | controls, is controlled by, or is under common control with You. For 80 | purposes of this definition, "control" means (a) the power, direct 81 | or indirect, to cause the direction or management of such entity, 82 | whether by contract or otherwise, or (b) ownership of more than 83 | fifty percent (50%) of the outstanding shares or beneficial 84 | ownership of such entity. 85 | 86 | 2. License Grants and Conditions 87 | -------------------------------- 88 | 89 | 2.1. Grants 90 | 91 | Each Contributor hereby grants You a world-wide, royalty-free, 92 | non-exclusive license: 93 | 94 | (a) under intellectual property rights (other than patent or trademark) 95 | Licensable by such Contributor to use, reproduce, make available, 96 | modify, display, perform, distribute, and otherwise exploit its 97 | Contributions, either on an unmodified basis, with Modifications, or 98 | as part of a Larger Work; and 99 | 100 | (b) under Patent Claims of such Contributor to make, use, sell, offer 101 | for sale, have made, import, and otherwise transfer either its 102 | Contributions or its Contributor Version. 103 | 104 | 2.2. Effective Date 105 | 106 | The licenses granted in Section 2.1 with respect to any Contribution 107 | become effective for each Contribution on the date the Contributor first 108 | distributes such Contribution. 109 | 110 | 2.3. Limitations on Grant Scope 111 | 112 | The licenses granted in this Section 2 are the only rights granted under 113 | this License. No additional rights or licenses will be implied from the 114 | distribution or licensing of Covered Software under this License. 115 | Notwithstanding Section 2.1(b) above, no patent license is granted by a 116 | Contributor: 117 | 118 | (a) for any code that a Contributor has removed from Covered Software; 119 | or 120 | 121 | (b) for infringements caused by: (i) Your and any other third party's 122 | modifications of Covered Software, or (ii) the combination of its 123 | Contributions with other software (except as part of its Contributor 124 | Version); or 125 | 126 | (c) under Patent Claims infringed by Covered Software in the absence of 127 | its Contributions. 128 | 129 | This License does not grant any rights in the trademarks, service marks, 130 | or logos of any Contributor (except as may be necessary to comply with 131 | the notice requirements in Section 3.4). 132 | 133 | 2.4. Subsequent Licenses 134 | 135 | No Contributor makes additional grants as a result of Your choice to 136 | distribute the Covered Software under a subsequent version of this 137 | License (see Section 10.2) or under the terms of a Secondary License (if 138 | permitted under the terms of Section 3.3). 139 | 140 | 2.5. Representation 141 | 142 | Each Contributor represents that the Contributor believes its 143 | Contributions are its original creation(s) or it has sufficient rights 144 | to grant the rights to its Contributions conveyed by this License. 145 | 146 | 2.6. Fair Use 147 | 148 | This License is not intended to limit any rights You have under 149 | applicable copyright doctrines of fair use, fair dealing, or other 150 | equivalents. 151 | 152 | 2.7. Conditions 153 | 154 | Sections 3.1, 3.2, 3.3, and 3.4 are conditions of the licenses granted 155 | in Section 2.1. 156 | 157 | 3. Responsibilities 158 | ------------------- 159 | 160 | 3.1. Distribution of Source Form 161 | 162 | All distribution of Covered Software in Source Code Form, including any 163 | Modifications that You create or to which You contribute, must be under 164 | the terms of this License. You must inform recipients that the Source 165 | Code Form of the Covered Software is governed by the terms of this 166 | License, and how they can obtain a copy of this License. You may not 167 | attempt to alter or restrict the recipients' rights in the Source Code 168 | Form. 169 | 170 | 3.2. Distribution of Executable Form 171 | 172 | If You distribute Covered Software in Executable Form then: 173 | 174 | (a) such Covered Software must also be made available in Source Code 175 | Form, as described in Section 3.1, and You must inform recipients of 176 | the Executable Form how they can obtain a copy of such Source Code 177 | Form by reasonable means in a timely manner, at a charge no more 178 | than the cost of distribution to the recipient; and 179 | 180 | (b) You may distribute such Executable Form under the terms of this 181 | License, or sublicense it under different terms, provided that the 182 | license for the Executable Form does not attempt to limit or alter 183 | the recipients' rights in the Source Code Form under this License. 184 | 185 | 3.3. Distribution of a Larger Work 186 | 187 | You may create and distribute a Larger Work under terms of Your choice, 188 | provided that You also comply with the requirements of this License for 189 | the Covered Software. If the Larger Work is a combination of Covered 190 | Software with a work governed by one or more Secondary Licenses, and the 191 | Covered Software is not Incompatible With Secondary Licenses, this 192 | License permits You to additionally distribute such Covered Software 193 | under the terms of such Secondary License(s), so that the recipient of 194 | the Larger Work may, at their option, further distribute the Covered 195 | Software under the terms of either this License or such Secondary 196 | License(s). 197 | 198 | 3.4. Notices 199 | 200 | You may not remove or alter the substance of any license notices 201 | (including copyright notices, patent notices, disclaimers of warranty, 202 | or limitations of liability) contained within the Source Code Form of 203 | the Covered Software, except that You may alter any license notices to 204 | the extent required to remedy known factual inaccuracies. 205 | 206 | 3.5. Application of Additional Terms 207 | 208 | You may choose to offer, and to charge a fee for, warranty, support, 209 | indemnity or liability obligations to one or more recipients of Covered 210 | Software. However, You may do so only on Your own behalf, and not on 211 | behalf of any Contributor. You must make it absolutely clear that any 212 | such warranty, support, indemnity, or liability obligation is offered by 213 | You alone, and You hereby agree to indemnify every Contributor for any 214 | liability incurred by such Contributor as a result of warranty, support, 215 | indemnity or liability terms You offer. You may include additional 216 | disclaimers of warranty and limitations of liability specific to any 217 | jurisdiction. 218 | 219 | 4. Inability to Comply Due to Statute or Regulation 220 | --------------------------------------------------- 221 | 222 | If it is impossible for You to comply with any of the terms of this 223 | License with respect to some or all of the Covered Software due to 224 | statute, judicial order, or regulation then You must: (a) comply with 225 | the terms of this License to the maximum extent possible; and (b) 226 | describe the limitations and the code they affect. Such description must 227 | be placed in a text file included with all distributions of the Covered 228 | Software under this License. Except to the extent prohibited by statute 229 | or regulation, such description must be sufficiently detailed for a 230 | recipient of ordinary skill to be able to understand it. 231 | 232 | 5. Termination 233 | -------------- 234 | 235 | 5.1. The rights granted under this License will terminate automatically 236 | if You fail to comply with any of its terms. However, if You become 237 | compliant, then the rights granted under this License from a particular 238 | Contributor are reinstated (a) provisionally, unless and until such 239 | Contributor explicitly and finally terminates Your grants, and (b) on an 240 | ongoing basis, if such Contributor fails to notify You of the 241 | non-compliance by some reasonable means prior to 60 days after You have 242 | come back into compliance. Moreover, Your grants from a particular 243 | Contributor are reinstated on an ongoing basis if such Contributor 244 | notifies You of the non-compliance by some reasonable means, this is the 245 | first time You have received notice of non-compliance with this License 246 | from such Contributor, and You become compliant prior to 30 days after 247 | Your receipt of the notice. 248 | 249 | 5.2. If You initiate litigation against any entity by asserting a patent 250 | infringement claim (excluding declaratory judgment actions, 251 | counter-claims, and cross-claims) alleging that a Contributor Version 252 | directly or indirectly infringes any patent, then the rights granted to 253 | You by any and all Contributors for the Covered Software under Section 254 | 2.1 of this License shall terminate. 255 | 256 | 5.3. In the event of termination under Sections 5.1 or 5.2 above, all 257 | end user license agreements (excluding distributors and resellers) which 258 | have been validly granted by You or Your distributors under this License 259 | prior to termination shall survive termination. 260 | 261 | ************************************************************************ 262 | * * 263 | * 6. Disclaimer of Warranty * 264 | * ------------------------- * 265 | * * 266 | * Covered Software is provided under this License on an "as is" * 267 | * basis, without warranty of any kind, either expressed, implied, or * 268 | * statutory, including, without limitation, warranties that the * 269 | * Covered Software is free of defects, merchantable, fit for a * 270 | * particular purpose or non-infringing. The entire risk as to the * 271 | * quality and performance of the Covered Software is with You. * 272 | * Should any Covered Software prove defective in any respect, You * 273 | * (not any Contributor) assume the cost of any necessary servicing, * 274 | * repair, or correction. This disclaimer of warranty constitutes an * 275 | * essential part of this License. No use of any Covered Software is * 276 | * authorized under this License except under this disclaimer. * 277 | * * 278 | ************************************************************************ 279 | 280 | ************************************************************************ 281 | * * 282 | * 7. Limitation of Liability * 283 | * -------------------------- * 284 | * * 285 | * Under no circumstances and under no legal theory, whether tort * 286 | * (including negligence), contract, or otherwise, shall any * 287 | * Contributor, or anyone who distributes Covered Software as * 288 | * permitted above, be liable to You for any direct, indirect, * 289 | * special, incidental, or consequential damages of any character * 290 | * including, without limitation, damages for lost profits, loss of * 291 | * goodwill, work stoppage, computer failure or malfunction, or any * 292 | * and all other commercial damages or losses, even if such party * 293 | * shall have been informed of the possibility of such damages. This * 294 | * limitation of liability shall not apply to liability for death or * 295 | * personal injury resulting from such party's negligence to the * 296 | * extent applicable law prohibits such limitation. Some * 297 | * jurisdictions do not allow the exclusion or limitation of * 298 | * incidental or consequential damages, so this exclusion and * 299 | * limitation may not apply to You. * 300 | * * 301 | ************************************************************************ 302 | 303 | 8. Litigation 304 | ------------- 305 | 306 | Any litigation relating to this License may be brought only in the 307 | courts of a jurisdiction where the defendant maintains its principal 308 | place of business and such litigation shall be governed by laws of that 309 | jurisdiction, without reference to its conflict-of-law provisions. 310 | Nothing in this Section shall prevent a party's ability to bring 311 | cross-claims or counter-claims. 312 | 313 | 9. Miscellaneous 314 | ---------------- 315 | 316 | This License represents the complete agreement concerning the subject 317 | matter hereof. If any provision of this License is held to be 318 | unenforceable, such provision shall be reformed only to the extent 319 | necessary to make it enforceable. Any law or regulation which provides 320 | that the language of a contract shall be construed against the drafter 321 | shall not be used to construe this License against a Contributor. 322 | 323 | 10. Versions of the License 324 | --------------------------- 325 | 326 | 10.1. New Versions 327 | 328 | Mozilla Foundation is the license steward. Except as provided in Section 329 | 10.3, no one other than the license steward has the right to modify or 330 | publish new versions of this License. Each version will be given a 331 | distinguishing version number. 332 | 333 | 10.2. Effect of New Versions 334 | 335 | You may distribute the Covered Software under the terms of the version 336 | of the License under which You originally received the Covered Software, 337 | or under the terms of any subsequent version published by the license 338 | steward. 339 | 340 | 10.3. Modified Versions 341 | 342 | If you create software not governed by this License, and you want to 343 | create a new license for such software, you may create and use a 344 | modified version of this License if you rename the license and remove 345 | any references to the name of the license steward (except to note that 346 | such modified license differs from this License). 347 | 348 | 10.4. Distributing Source Code Form that is Incompatible With Secondary 349 | Licenses 350 | 351 | If You choose to distribute Source Code Form that is Incompatible With 352 | Secondary Licenses under the terms of this version of the License, the 353 | notice described in Exhibit B of this License must be attached. 354 | 355 | Exhibit A - Source Code Form License Notice 356 | ------------------------------------------- 357 | 358 | This Source Code Form is subject to the terms of the Mozilla Public 359 | License, v. 2.0. If a copy of the MPL was not distributed with this 360 | file, You can obtain one at http://mozilla.org/MPL/2.0/. 361 | 362 | If it is not possible or desirable to put the notice in a particular 363 | file, then You may include the notice in a location (such as a LICENSE 364 | file in a relevant directory) where a recipient would be likely to look 365 | for such a notice. 366 | 367 | You may add additional accurate notices of copyright ownership. 368 | 369 | Exhibit B - "Incompatible With Secondary Licenses" Notice 370 | --------------------------------------------------------- 371 | 372 | This Source Code Form is "Incompatible With Secondary Licenses", as 373 | defined by the Mozilla Public License, v. 2.0. 374 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # OneCache 2 | OneCache is a distributed key/value store that is accessible via the memcached 3 | protocol. OneCache nodes natively support clustering and best effort 4 | replication. Adding nodes into the cluster effectively increases the in-memory 5 | storage capability of the cache. Replication of keys minimizes the impact of a 6 | single node failure as requests will transparently be forwarded to the new key 7 | leader. 8 | 9 | OneCache's architecture is built around the concept of a one-hop distributed 10 | hash table. Many DHT algorithms are designed for a very large number of 11 | members, such as bitorrent DHTs, and trade lookup efficiency for per-node 12 | storage and communication. However, OneCache is built for a much smaller 13 | cluster size (<100) and as such, every node holds the full state of the 14 | underlying consistent hash ring. This is advantageous because every node is 15 | capable of forwarding requests to the correct node, which is why it receives 16 | the term one-hop. 17 | 18 | OneCache's name then becomes a play on its one-hop DHT architecture and memcached protocol 19 | support. 20 | 21 | Download pre-compiled binaries by browsing the [releases.](https://github.com/dadgar/onecache/releases) 22 | 23 | See a micro-benchmark comparing OneCache and memcached [here.](BENCHMARK.md) 24 | 25 | ## Usage 26 | ### Single Node 27 | To start a single OneCache node listening on memcached's default port, simply 28 | run: 29 | ``` 30 | $> ./onecache 31 | ``` 32 | 33 | ### Clustered 34 | To start a cluster of OneCache nodes on a single machine (for ease of getting 35 | started) we must bind the nodes to non-default ports to avoid collision. If 36 | clustering on different machines, the default ports can be used. 37 | 38 | Start the first node in one terminal: 39 | ``` 40 | $> ./onecache -gossip_port=7946 -port=11211 41 | ``` 42 | 43 | In another, start the second and have it join. The output will show it added a 44 | peer node. 45 | ``` 46 | $> ./onecache -gossip_port=7947 -port=11212 -join="127.0.0.1:7946" 47 | 2016/01/18 12:04:30 [INFO] onecache: added peer node &{c3f48bd0-4106-43d4-9044-15cf552037f0 192.168.1.10 11211} 48 | ``` 49 | 50 | ## Usage 51 | Standard memcached libraries can be used, which make OneCache an easy drop in 52 | replacement for memcached as code doesn't have to be changed. OneCache is also 53 | accessible via the terminal: 54 | 55 | ``` 56 | $> telnet 127.0.0.1 11211 57 | Trying 127.0.0.1... 58 | Connected to localhost. 59 | Escape character is '^]'. 60 | set foo 0 0 12 61 | hello, world 62 | STORED 63 | ``` 64 | -------------------------------------------------------------------------------- /main.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "flag" 5 | "fmt" 6 | "log" 7 | "net/http" 8 | "os" 9 | "os/signal" 10 | "strings" 11 | "time" 12 | 13 | "github.com/dadgar/onecache/onecache" 14 | 15 | _ "net/http/pprof" 16 | ) 17 | 18 | var ( 19 | defaultConfig = onecache.DefaultNodeConfig() 20 | join = flag.String("join", "", "A comma seperated list of peers (ip:port/dns) to join") 21 | retry_join = flag.String("retry_join", "", "A comma seperated list of peers (ip:port/dns) to join with retries enabled") 22 | retry_interval = flag.Duration("retry_interval", 30*time.Second, "The interval to try retrying a join if there was an error") 23 | port = flag.Int("port", defaultConfig.RpcPort, "Port to bind to for memcache connections") 24 | gossipPort = flag.Int("gossip_port", defaultConfig.SerfConfig.MemberlistConfig.BindPort, "Port to use for membership gossiping") 25 | addr = flag.String("addr", "", "Address to bind to") 26 | replicas = flag.Int("replicas", defaultConfig.Replicas, "Number of replicas") 27 | maxMemory = flag.Int64("max_memory", defaultConfig.MaxMemoryUsage, "Maximum memory to use (in bytes)") 28 | pprofEnabled = flag.Bool("pprof", false, "Enable pprof for debugging") 29 | pprofPort = flag.Int("pprof_port", 6060, "Pprof http port") 30 | ) 31 | 32 | func main() { 33 | flag.Parse() 34 | config := buildConfig() 35 | logger := log.New(config.LogOutput, "", log.LstdFlags) 36 | node, err := onecache.Create(config) 37 | if err != nil { 38 | logger.Printf("[ERROR] onecache: could not create node: %v\n", err) 39 | os.Exit(1) 40 | } 41 | 42 | // Handle interupts. 43 | c := make(chan os.Signal, 1) 44 | signal.Notify(c, os.Interrupt) 45 | 46 | node.Start() 47 | 48 | if *join != "" { 49 | peerNodes := strings.Split(*join, ",") 50 | if err := node.Join(peerNodes); err != nil { 51 | logger.Printf("[ERROR] onecache: could not join peers %v: %v", peerNodes, err) 52 | node.Exit() 53 | return 54 | } 55 | } 56 | 57 | if *retry_join != "" { 58 | var retry *time.Timer 59 | retryNodes := strings.Split(*retry_join, ",") 60 | retry = time.AfterFunc(0, func() { 61 | if err := node.Join(retryNodes); err != nil { 62 | logger.Printf("[DEBUG] onecache: could not join peers %v: %v\n", retryNodes, err) 63 | logger.Printf("[DEBUG] onecache: retrying in %v", *retry_interval) 64 | retry.Reset(*retry_interval) 65 | } else { 66 | retry.Stop() 67 | } 68 | }) 69 | } 70 | 71 | if *pprofEnabled { 72 | go func() { 73 | logger.Println(http.ListenAndServe(fmt.Sprintf(":%d", *pprofPort), nil)) 74 | }() 75 | } 76 | 77 | // Exit the node upon an interupt. 78 | for range c { 79 | node.Exit() 80 | return 81 | } 82 | } 83 | 84 | func buildConfig() *onecache.NodeConfig { 85 | config := onecache.DefaultNodeConfig() 86 | config.LogOutput = os.Stderr 87 | 88 | if *port != 0 { 89 | config.RpcPort = *port 90 | } 91 | 92 | if *gossipPort != 0 { 93 | config.SerfConfig.MemberlistConfig.BindPort = *gossipPort 94 | } 95 | 96 | if *addr != "" { 97 | config.SerfConfig.MemberlistConfig.BindAddr = *addr 98 | } 99 | 100 | if *replicas != 0 { 101 | config.Replicas = *replicas 102 | } 103 | 104 | if *maxMemory != -1 { 105 | config.MaxMemoryUsage = *maxMemory 106 | } 107 | 108 | return config 109 | } 110 | -------------------------------------------------------------------------------- /onecache/config.go: -------------------------------------------------------------------------------- 1 | package onecache 2 | 3 | import ( 4 | "io" 5 | 6 | "github.com/hashicorp/serf/serf" 7 | ) 8 | 9 | // NodeConfig is the configuration for a DHT Node instance. 10 | type NodeConfig struct { 11 | // The number of key/value replicas. 12 | Replicas int 13 | 14 | // The port to accept RPC connections on. 15 | RpcPort int 16 | 17 | // The space alloted to store key/values in bytes. 18 | MaxMemoryUsage int64 19 | 20 | // Routes log output to the writer. If nil, output will be sent to stderr. 21 | LogOutput io.Writer 22 | 23 | // Serf configuration used to fine tune Node discovery, define whether the DHT is running in a LAN or WAN and more. 24 | // The name does not need to be set as it will be overriden with a UUID. 25 | SerfConfig *serf.Config 26 | } 27 | 28 | // DefaultNodeConfig returns a reasonable default to use. 29 | func DefaultNodeConfig() *NodeConfig { 30 | return &NodeConfig{ 31 | Replicas: 2, 32 | RpcPort: 11211, 33 | SerfConfig: serf.DefaultConfig(), 34 | MaxMemoryUsage: 1 << 30, // 1 GB 35 | } 36 | } 37 | -------------------------------------------------------------------------------- /onecache/memcache.go: -------------------------------------------------------------------------------- 1 | package onecache 2 | 3 | import ( 4 | "bufio" 5 | "fmt" 6 | "io" 7 | "net" 8 | "strings" 9 | ) 10 | 11 | var handlers map[string]commandHandler 12 | 13 | // getHandlers returns the memcache command handlers. 14 | func getHandlers(n *Node) map[string]commandHandler { 15 | if len(handlers) == 0 { 16 | handlers = map[string]commandHandler{ 17 | "add": &addHandler{n: n}, 18 | "append": &appendHandler{n: n, back: false}, 19 | "cas": &casHandler{n: n}, 20 | "decr": &incrementHandler{n: n, negative: true}, 21 | "delete": &deleteHandler{n: n}, 22 | "get": &getHandler{n: n, cas: false}, 23 | "gets": &getHandler{n: n, cas: true}, 24 | "incr": &incrementHandler{n: n, negative: false}, 25 | "prepend": &appendHandler{n: n, back: true}, 26 | "replace": &replaceHandler{n: n}, 27 | "set": &setHandler{n: n}, 28 | "touch": &touchHandler{n: n}, 29 | } 30 | } 31 | 32 | return handlers 33 | } 34 | 35 | // wrappedByteReader is a reader that first returns the wrapped byte and 36 | // subsequently reads from the reader. 37 | type wrappedByteReader struct { 38 | first byte 39 | r io.Reader 40 | sent bool 41 | } 42 | 43 | // newWrappedByteReader returns a reader that first returns the passed byte and 44 | // then reads from the reader. 45 | func newWrappedByteReader(b byte, r io.Reader) *wrappedByteReader { 46 | return &wrappedByteReader{ 47 | first: b, 48 | r: r, 49 | sent: false, 50 | } 51 | } 52 | 53 | // Read first reads the wrapped byte and then reads from the reader. 54 | func (w *wrappedByteReader) Read(p []byte) (n int, err error) { 55 | if len(p) == 0 { 56 | return 0, fmt.Errorf("passed empty byte slice") 57 | } 58 | 59 | if !w.sent { 60 | w.sent = true 61 | p[0] = w.first 62 | n, err := w.r.Read(p[1:]) 63 | n++ 64 | return n, err 65 | } 66 | 67 | return w.r.Read(p) 68 | } 69 | 70 | // handleMemcacheClient parses and responds to memcache queries. It is passed 71 | // the first byte of the stream. 72 | func (n *Node) handleMemcacheClient(conn net.Conn, first byte) { 73 | // Setup the handlers. 74 | handlers := getHandlers(n) 75 | 76 | // Create a buffered connection 77 | r := bufio.NewReader(newWrappedByteReader(first, conn)) 78 | w := bufio.NewWriter(conn) 79 | buf := bufio.NewReadWriter(r, w) 80 | 81 | for { 82 | s, err := buf.ReadString('\n') 83 | if err != nil { 84 | if err := sendClientError(buf, "invalid input: couldn't read till end of line"); err != nil { 85 | buf.Flush() 86 | conn.Close() 87 | return 88 | } 89 | 90 | buf.Flush() 91 | continue 92 | } 93 | 94 | s = strings.TrimSuffix(s, "\r\n") 95 | input := strings.Split(s, " ") 96 | if len(input) == 1 && input[0] == "" { 97 | if err := sendClientError(buf, "invalid input: sent no commands"); err != nil { 98 | buf.Flush() 99 | conn.Close() 100 | return 101 | } 102 | 103 | buf.Flush() 104 | continue 105 | } 106 | 107 | cmd := input[0] 108 | handler, ok := handlers[cmd] 109 | if !ok { 110 | if err := sendError(buf, fmt.Sprintf("unknown command %v", cmd)); err != nil { 111 | buf.Flush() 112 | conn.Close() 113 | return 114 | } 115 | 116 | buf.Flush() 117 | continue 118 | } 119 | 120 | if err := handler.parse(input[1:]); err != nil { 121 | if err := sendClientError(buf, err.Error()); err != nil { 122 | buf.Flush() 123 | conn.Close() 124 | return 125 | } 126 | 127 | buf.Flush() 128 | continue 129 | } 130 | 131 | b := handler.readBytes() 132 | var data []byte 133 | if b > 0 { 134 | data = make([]byte, b) 135 | if _, err := buf.Read(data); err != nil { 136 | if err := sendClientError(buf, err.Error()); err != nil { 137 | buf.Flush() 138 | conn.Close() 139 | return 140 | } 141 | buf.Flush() 142 | continue 143 | } 144 | } 145 | 146 | if err := handler.run(buf, data); err != nil { 147 | if err := sendServerError(buf, err.Error()); err != nil { 148 | buf.Flush() 149 | conn.Close() 150 | return 151 | } 152 | buf.Flush() 153 | conn.Close() 154 | return 155 | } 156 | 157 | buf.Flush() 158 | 159 | // Strip the remaining terminal. 160 | if b > 0 { 161 | if _, err := buf.Discard(2); err != nil { 162 | if err := sendServerError(buf, err.Error()); err != nil { 163 | buf.Flush() 164 | conn.Close() 165 | return 166 | } 167 | buf.Flush() 168 | conn.Close() 169 | return 170 | } 171 | } 172 | } 173 | } 174 | 175 | // Helpers to send error messages back to the client. 176 | func sendError(w io.Writer, msg string) error { 177 | return sendGenericError(w, "ERROR", msg) 178 | } 179 | 180 | func sendClientError(w io.Writer, msg string) error { 181 | return sendGenericError(w, "CLIENT_ERROR", msg) 182 | } 183 | 184 | func sendServerError(w io.Writer, msg string) error { 185 | return sendGenericError(w, "SERVER_ERROR", msg) 186 | } 187 | 188 | func sendGenericError(w io.Writer, prefix, msg string) error { 189 | errMsg := fmt.Sprintf("%v %v\r\n", prefix, msg) 190 | if _, err := io.WriteString(w, errMsg); err != nil { 191 | return err 192 | } 193 | return nil 194 | } 195 | 196 | // sendMsg writes the message, appending the correct delimiter. 197 | func sendMsg(w io.Writer, msg string) error { 198 | delimited := fmt.Sprintf("%v\r\n", msg) 199 | if _, err := io.WriteString(w, delimited); err != nil { 200 | return err 201 | } 202 | return nil 203 | } 204 | -------------------------------------------------------------------------------- /onecache/memcache_handlers.go: -------------------------------------------------------------------------------- 1 | package onecache 2 | 3 | import ( 4 | "fmt" 5 | "io" 6 | "strconv" 7 | "time" 8 | 9 | "github.com/dadgar/onecache/ttlstore" 10 | ) 11 | 12 | const ( 13 | noreply = "noreply" 14 | 15 | expectedSetInput = " [noreply]\r\n" 16 | expectedCasInput = " [noreply]\r\n" 17 | expectedIncrInput = " [noreply]\r\n" 18 | expectedGetInput = "*\r\n" 19 | expectedDeleteInput = " [noreply]\r\n" 20 | expectedTouchInput = " [noreply]\r\n" 21 | 22 | // The maximum value in which expirations are considered as seconds from now. 23 | secondMax = 60 * 60 * 24 * 30 24 | ) 25 | 26 | // commandHandler is the interface to handle a memcache command. 27 | type commandHandler interface { 28 | // Returns an error if the input is not properly formatted. 29 | parse(flags []string) error 30 | 31 | // Returns how many bytes need to be read off the connection before calling Run(). -1 indicates that no additional 32 | // reading is necessary. 33 | readBytes() int64 34 | 35 | // Run the command returning the results to w. Must be called after Parse. Required input will be passed through data. 36 | // If an unrecoverable error occurs, error will not be nil and the connection to the client will be closed. 37 | run(w io.Writer, data []byte) error 38 | } 39 | 40 | func convertMemcacheExpToUnix(memExp int32) int64 { 41 | exp := int64(memExp) 42 | if memExp == 0 { 43 | return 0 44 | } else if exp < secondMax { 45 | return time.Now().Add(time.Duration(exp) * time.Second).Unix() 46 | } 47 | 48 | return exp 49 | } 50 | 51 | type setHandler struct { 52 | n *Node 53 | args StorageArgs 54 | size int64 55 | noReply bool 56 | } 57 | 58 | func (h *setHandler) parse(flags []string) error { 59 | l := len(flags) 60 | if l != 4 && l != 5 { 61 | return fmt.Errorf("received an unexpected number of flags (%d); expect %v", l, expectedSetInput) 62 | } 63 | 64 | h.args.Key = flags[0] 65 | 66 | flag, err := strconv.ParseInt(flags[1], 10, 32) 67 | if err != nil { 68 | return fmt.Errorf("could not parse flag field: %v", err) 69 | } 70 | h.args.Flags = int32(flag) 71 | 72 | exp, err := strconv.ParseInt(flags[2], 10, 32) 73 | if err != nil { 74 | return fmt.Errorf("could not parse expiration field: %v", err) 75 | return err 76 | } 77 | h.args.Exp = convertMemcacheExpToUnix(int32(exp)) 78 | 79 | bytes, err := strconv.ParseInt(flags[3], 10, 64) 80 | if err != nil { 81 | return fmt.Errorf("could not parse bytes field: %v", err) 82 | } 83 | h.size = bytes 84 | 85 | if l == 5 { 86 | if flags[4] != noreply { 87 | return fmt.Errorf("was expecting %v; got %v", noreply, flags[4]) 88 | } 89 | 90 | h.noReply = true 91 | } 92 | 93 | return nil 94 | } 95 | 96 | func (h *setHandler) readBytes() int64 { 97 | return h.size 98 | } 99 | 100 | func (h *setHandler) run(w io.Writer, data []byte) error { 101 | var resp struct{} 102 | h.args.Value = data 103 | err := h.n.server.Set(h.args, &resp) 104 | if h.noReply { 105 | return nil 106 | } 107 | 108 | if err != nil { 109 | return sendMsg(w, "NOT_STORED") 110 | } 111 | 112 | return sendMsg(w, "STORED") 113 | } 114 | 115 | type casHandler struct { 116 | n *Node 117 | args StorageArgs 118 | size int64 119 | noReply bool 120 | } 121 | 122 | func (h *casHandler) parse(flags []string) error { 123 | l := len(flags) 124 | if l != 5 && l != 6 { 125 | return fmt.Errorf("received an unexpected number of flags (%d); expect %v", l, expectedCasInput) 126 | } 127 | 128 | h.args.Key = flags[0] 129 | 130 | flag, err := strconv.ParseInt(flags[1], 10, 32) 131 | if err != nil { 132 | return fmt.Errorf("could not parse flag field: %v", err) 133 | } 134 | h.args.Flags = int32(flag) 135 | 136 | exp, err := strconv.ParseInt(flags[2], 10, 32) 137 | if err != nil { 138 | return fmt.Errorf("could not parse expiration field: %v", err) 139 | return err 140 | } 141 | h.args.Exp = convertMemcacheExpToUnix(int32(exp)) 142 | 143 | bytes, err := strconv.ParseInt(flags[3], 10, 64) 144 | if err != nil { 145 | return fmt.Errorf("could not parse bytes field: %v", err) 146 | } 147 | h.size = bytes 148 | 149 | cas, err := strconv.ParseInt(flags[4], 10, 64) 150 | if err != nil { 151 | return fmt.Errorf("could not parse cas field: %v", err) 152 | } 153 | h.args.Cas = cas 154 | 155 | if l == 6 { 156 | if flags[5] != noreply { 157 | return fmt.Errorf("was expecting %v; got %v", noreply, flags[5]) 158 | } 159 | 160 | h.noReply = true 161 | } 162 | 163 | return nil 164 | } 165 | 166 | func (h *casHandler) readBytes() int64 { 167 | return h.size 168 | } 169 | 170 | func (h *casHandler) run(w io.Writer, data []byte) error { 171 | h.args.Value = data 172 | var resp struct{} 173 | err := h.n.server.Cas(h.args, &resp) 174 | if h.noReply { 175 | return nil 176 | } 177 | 178 | if err != nil { 179 | return sendMsg(w, "NOT_STORED") 180 | } 181 | 182 | return sendMsg(w, "STORED") 183 | } 184 | 185 | type addHandler struct { 186 | n *Node 187 | args StorageArgs 188 | size int64 189 | noReply bool 190 | } 191 | 192 | func (h *addHandler) parse(flags []string) error { 193 | l := len(flags) 194 | if l != 4 && l != 5 { 195 | return fmt.Errorf("received an unexpected number of flags (%d); expect %v", l, expectedSetInput) 196 | } 197 | 198 | h.args.Key = flags[0] 199 | 200 | flag, err := strconv.ParseInt(flags[1], 10, 32) 201 | if err != nil { 202 | return fmt.Errorf("could not parse flag field: %v", err) 203 | } 204 | h.args.Flags = int32(flag) 205 | 206 | exp, err := strconv.ParseInt(flags[2], 10, 32) 207 | if err != nil { 208 | return fmt.Errorf("could not parse expiration field: %v", err) 209 | return err 210 | } 211 | h.args.Exp = convertMemcacheExpToUnix(int32(exp)) 212 | 213 | bytes, err := strconv.ParseInt(flags[3], 10, 64) 214 | if err != nil { 215 | return fmt.Errorf("could not parse bytes field: %v", err) 216 | } 217 | h.size = bytes 218 | 219 | if l == 5 { 220 | if flags[4] != noreply { 221 | return fmt.Errorf("was expecting %v; got %v", noreply, flags[4]) 222 | } 223 | 224 | h.noReply = true 225 | } 226 | 227 | return nil 228 | } 229 | 230 | func (h *addHandler) readBytes() int64 { 231 | return h.size 232 | } 233 | 234 | func (h *addHandler) run(w io.Writer, data []byte) error { 235 | var resp struct{} 236 | h.args.Value = data 237 | err := h.n.server.Add(h.args, &resp) 238 | if h.noReply { 239 | return nil 240 | } 241 | 242 | if err != nil { 243 | return sendMsg(w, "NOT_STORED") 244 | } 245 | 246 | return sendMsg(w, "STORED") 247 | } 248 | 249 | type replaceHandler struct { 250 | n *Node 251 | args StorageArgs 252 | size int64 253 | noReply bool 254 | } 255 | 256 | func (h *replaceHandler) parse(flags []string) error { 257 | l := len(flags) 258 | if l != 4 && l != 5 { 259 | return fmt.Errorf("received an unexpected number of flags (%d); expect %v", l, expectedSetInput) 260 | } 261 | 262 | h.args.Key = flags[0] 263 | 264 | flag, err := strconv.ParseInt(flags[1], 10, 32) 265 | if err != nil { 266 | return fmt.Errorf("could not parse flag field: %v", err) 267 | } 268 | h.args.Flags = int32(flag) 269 | 270 | exp, err := strconv.ParseInt(flags[2], 10, 32) 271 | if err != nil { 272 | return fmt.Errorf("could not parse expiration field: %v", err) 273 | return err 274 | } 275 | h.args.Exp = convertMemcacheExpToUnix(int32(exp)) 276 | 277 | bytes, err := strconv.ParseInt(flags[3], 10, 64) 278 | if err != nil { 279 | return fmt.Errorf("could not parse bytes field: %v", err) 280 | } 281 | h.size = bytes 282 | 283 | if l == 5 { 284 | if flags[4] != noreply { 285 | return fmt.Errorf("was expecting %v; got %v", noreply, flags[4]) 286 | } 287 | 288 | h.noReply = true 289 | } 290 | 291 | return nil 292 | } 293 | 294 | func (h *replaceHandler) readBytes() int64 { 295 | return h.size 296 | } 297 | 298 | func (h *replaceHandler) run(w io.Writer, data []byte) error { 299 | var resp struct{} 300 | h.args.Value = data 301 | err := h.n.server.Replace(h.args, &resp) 302 | if h.noReply { 303 | return nil 304 | } 305 | 306 | if err != nil { 307 | return sendMsg(w, "NOT_STORED") 308 | } 309 | 310 | return sendMsg(w, "STORED") 311 | } 312 | 313 | type appendHandler struct { 314 | n *Node 315 | back bool 316 | args StorageArgs 317 | size int64 318 | noReply bool 319 | } 320 | 321 | func (h *appendHandler) parse(flags []string) error { 322 | l := len(flags) 323 | if l != 4 && l != 5 { 324 | return fmt.Errorf("received an unexpected number of flags (%d); expect %v", l, expectedSetInput) 325 | } 326 | 327 | h.args.Key = flags[0] 328 | 329 | bytes, err := strconv.ParseInt(flags[3], 10, 64) 330 | if err != nil { 331 | return fmt.Errorf("could not parse bytes field: %v", err) 332 | } 333 | h.size = bytes 334 | 335 | if l == 5 { 336 | if flags[4] != noreply { 337 | return fmt.Errorf("was expecting %v; got %v", noreply, flags[4]) 338 | } 339 | 340 | h.noReply = true 341 | } 342 | 343 | return nil 344 | } 345 | 346 | func (h *appendHandler) readBytes() int64 { 347 | return h.size 348 | } 349 | 350 | func (h *appendHandler) run(w io.Writer, data []byte) error { 351 | var resp struct{} 352 | var err error 353 | h.args.Value = data 354 | if h.back { 355 | err = h.n.server.Prepend(h.args, &resp) 356 | } else { 357 | err = h.n.server.Append(h.args, &resp) 358 | } 359 | 360 | if h.noReply { 361 | return nil 362 | } 363 | 364 | if err != nil { 365 | return sendMsg(w, "NOT_STORED") 366 | } 367 | 368 | return sendMsg(w, "STORED") 369 | } 370 | 371 | type incrementHandler struct { 372 | n *Node 373 | negative bool 374 | args StorageArgs 375 | noReply bool 376 | } 377 | 378 | func (h *incrementHandler) parse(flags []string) error { 379 | l := len(flags) 380 | if l != 2 && l != 3 { 381 | return fmt.Errorf("received an unexpected number of flags (%d); expect %v", l, expectedIncrInput) 382 | } 383 | 384 | h.args.Key = flags[0] 385 | 386 | delta, err := strconv.ParseInt(flags[1], 10, 64) 387 | if err != nil { 388 | return fmt.Errorf("could not parse value: %v", err) 389 | } 390 | h.args.Delta = delta 391 | 392 | if l == 3 { 393 | if flags[2] != noreply { 394 | return fmt.Errorf("was expecting %v; got %v", noreply, flags[2]) 395 | } 396 | 397 | h.noReply = true 398 | } 399 | 400 | return nil 401 | } 402 | 403 | func (h *incrementHandler) readBytes() int64 { 404 | return 0 405 | } 406 | 407 | func (h *incrementHandler) run(w io.Writer, data []byte) error { 408 | var resp int64 409 | var err error 410 | if h.negative { 411 | err = h.n.server.Decrement(h.args, &resp) 412 | } else { 413 | err = h.n.server.Increment(h.args, &resp) 414 | } 415 | 416 | if h.noReply { 417 | return nil 418 | } 419 | 420 | if err != nil { 421 | return sendMsg(w, "NOT_FOUND") 422 | } 423 | 424 | return sendMsg(w, fmt.Sprintf("%d", resp)) 425 | } 426 | 427 | type getHandler struct { 428 | n *Node 429 | cas bool 430 | keys []string 431 | } 432 | 433 | func (h *getHandler) parse(flags []string) error { 434 | l := len(flags) 435 | if l == 0 { 436 | return fmt.Errorf("received an unexpected number of flags (%d); expect %v", l, expectedGetInput) 437 | } 438 | 439 | h.keys = flags 440 | return nil 441 | } 442 | 443 | func (h *getHandler) readBytes() int64 { 444 | return -1 445 | } 446 | 447 | func (h *getHandler) run(w io.Writer, data []byte) error { 448 | var resp ttlstore.KeyData 449 | var args StorageArgs 450 | for _, k := range h.keys { 451 | args.Key = k 452 | err := h.n.server.Get(args, &resp) 453 | if err != nil { 454 | continue 455 | } 456 | 457 | // Write results out. 458 | var header string 459 | if h.cas { 460 | header = fmt.Sprintf("VALUE %v %v %v %v", k, resp.Flags, resp.Size, resp.Cas) 461 | } else { 462 | header = fmt.Sprintf("VALUE %v %v %v", k, resp.Flags, resp.Size) 463 | } 464 | 465 | if err := sendMsg(w, header); err != nil { 466 | return err 467 | } 468 | 469 | if _, err := w.Write(resp.Data); err != nil { 470 | return err 471 | } 472 | 473 | if err := sendMsg(w, ""); err != nil { 474 | return err 475 | } 476 | } 477 | 478 | return sendMsg(w, "END") 479 | } 480 | 481 | type deleteHandler struct { 482 | n *Node 483 | key string 484 | noReply bool 485 | } 486 | 487 | func (h *deleteHandler) parse(flags []string) error { 488 | l := len(flags) 489 | if l != 1 && l != 2 { 490 | return fmt.Errorf("received an unexpected number of flags (%d); expect %v", l, expectedDeleteInput) 491 | } 492 | 493 | h.key = flags[0] 494 | 495 | if l == 2 { 496 | if flags[1] != noreply { 497 | return fmt.Errorf("was expecting %v; got %v", noreply, flags[1]) 498 | } 499 | 500 | h.noReply = true 501 | } 502 | 503 | return nil 504 | } 505 | 506 | func (h *deleteHandler) readBytes() int64 { 507 | return -1 508 | } 509 | 510 | func (h *deleteHandler) run(w io.Writer, data []byte) error { 511 | var args StorageArgs 512 | args.Key = h.key 513 | var resp struct{} 514 | err := h.n.server.Remove(args, &resp) 515 | if h.noReply { 516 | return nil 517 | } 518 | 519 | // TODO better error checking 520 | if err != nil { 521 | return sendMsg(w, "NOT_FOUND") 522 | } 523 | 524 | return sendMsg(w, "DELETED") 525 | } 526 | 527 | type touchHandler struct { 528 | n *Node 529 | args StorageArgs 530 | noReply bool 531 | } 532 | 533 | func (h *touchHandler) parse(flags []string) error { 534 | l := len(flags) 535 | if l != 2 && l != 3 { 536 | return fmt.Errorf("received an unexpected number of flags (%d); expect %v", l, expectedTouchInput) 537 | } 538 | 539 | h.args.Key = flags[0] 540 | 541 | exp, err := strconv.ParseInt(flags[1], 10, 32) 542 | if err != nil { 543 | return fmt.Errorf("could not parse expiration field: %v", err) 544 | return err 545 | } 546 | h.args.Exp = convertMemcacheExpToUnix(int32(exp)) 547 | 548 | if l == 3 { 549 | if flags[2] != noreply { 550 | return fmt.Errorf("was expecting %v; got %v", noreply, flags[2]) 551 | } 552 | 553 | h.noReply = true 554 | } 555 | 556 | return nil 557 | } 558 | 559 | func (h *touchHandler) readBytes() int64 { 560 | return -1 561 | } 562 | 563 | func (h *touchHandler) run(w io.Writer, data []byte) error { 564 | var resp struct{} 565 | err := h.n.server.Touch(h.args, &resp) 566 | if h.noReply { 567 | return nil 568 | } 569 | 570 | // TODO better error checking 571 | if err != nil { 572 | return sendMsg(w, "NOT_FOUND") 573 | } 574 | 575 | return sendMsg(w, "TOUCHED") 576 | return nil 577 | } 578 | -------------------------------------------------------------------------------- /onecache/memcache_test.go: -------------------------------------------------------------------------------- 1 | package onecache 2 | 3 | import ( 4 | "bufio" 5 | "fmt" 6 | "math/rand" 7 | "net" 8 | "testing" 9 | ) 10 | 11 | func BenchmarkMemcacheSetGet(b *testing.B) { 12 | n, err := defaultTestNode() 13 | if err != nil { 14 | b.Fatalf("defaultTestNode() failed: %v", err) 15 | } 16 | defer tearDownNodes([]*Node{n}) 17 | 18 | b.ResetTimer() 19 | b.RunParallel(func(pb *testing.PB) { 20 | conn, err := net.Dial("tcp", n.listener.Addr().String()) 21 | if err != nil { 22 | b.Fatalf("couldn't make connection to node: %v", err) 23 | } 24 | defer conn.Close() 25 | buf := bufio.NewReader(conn) 26 | 27 | for pb.Next() { 28 | r := rand.Int63() 29 | kv := fmt.Sprintf("%d", r%10000) 30 | 31 | if r%2 == 0 { 32 | fmt.Fprintf(conn, "GET %v\r\n", kv) 33 | if _, _, err := buf.ReadLine(); err != nil { 34 | b.Errorf("couldn't read response: %v", err) 35 | } 36 | } else { 37 | f := r % 10000 38 | ttl := r % 100 39 | s := len(kv) 40 | fmt.Fprintf(conn, "SET %v %d %d %d\r\n%v\r\n", kv, f, ttl, s, kv) 41 | if _, _, err := buf.ReadLine(); err != nil { 42 | b.Errorf("couldn't read response: %v", err) 43 | } 44 | } 45 | } 46 | }) 47 | } 48 | 49 | func BenchmarkMemcacheSetGetReplicated(b *testing.B) { 50 | nodes, err := connectedTestNodes(2, 2) 51 | defer tearDownNodes(nodes) 52 | if err != nil { 53 | b.Fatalf("connectedTestNodes(2, 2) failed: %v", err) 54 | } 55 | 56 | b.ResetTimer() 57 | b.RunParallel(func(pb *testing.PB) { 58 | conn1, err := net.Dial("tcp", nodes[0].listener.Addr().String()) 59 | if err != nil { 60 | b.Fatalf("couldn't make connection to node: %v", err) 61 | } 62 | defer conn1.Close() 63 | buf1 := bufio.NewReader(conn1) 64 | 65 | conn2, err := net.Dial("tcp", nodes[1].listener.Addr().String()) 66 | if err != nil { 67 | b.Fatalf("couldn't make connection to node: %v", err) 68 | } 69 | defer conn2.Close() 70 | buf2 := bufio.NewReader(conn2) 71 | 72 | var conn net.Conn 73 | var buf *bufio.Reader 74 | for pb.Next() { 75 | r := rand.Int63() 76 | r2 := rand.Int() 77 | kv := fmt.Sprintf("%d", r%10000) 78 | 79 | if r2%2 == 0 { 80 | conn = conn1 81 | buf = buf1 82 | } else { 83 | conn = conn2 84 | buf = buf2 85 | } 86 | 87 | if r%2 == 0 { 88 | fmt.Fprintf(conn, "GET %v\r\n", kv) 89 | if _, _, err := buf.ReadLine(); err != nil { 90 | b.Errorf("couldn't read response: %v", err) 91 | } 92 | } else { 93 | f := r % 10000 94 | ttl := r % 100 95 | s := len(kv) 96 | fmt.Fprintf(conn, "SET %v %d %d %d\r\n%v\r\n", kv, f, ttl, s, kv) 97 | if _, _, err := buf.ReadLine(); err != nil { 98 | b.Errorf("couldn't read response: %v", err) 99 | } 100 | } 101 | } 102 | }) 103 | } 104 | -------------------------------------------------------------------------------- /onecache/onecache.go: -------------------------------------------------------------------------------- 1 | package onecache 2 | 3 | import ( 4 | "errors" 5 | "fmt" 6 | "log" 7 | "net" 8 | "strconv" 9 | "sync" 10 | "time" 11 | 12 | "os" 13 | 14 | "github.com/dadgar/onecache/ttlstore" 15 | "github.com/hashicorp/serf/serf" 16 | "github.com/pborman/uuid" 17 | ) 18 | 19 | const ( 20 | // externalRpcTag is used to create a serf tag to broadcast the port being 21 | // used for RPC. 22 | externalRpcTag = "externalRpcTag" 23 | ) 24 | 25 | // peerNode holds the required information for sharding and forwarding RPCs to a 26 | // peer node in the ring. 27 | type peerNode struct { 28 | name string 29 | addr string 30 | port int 31 | } 32 | 33 | // Node holds all the necessary information to store data, handle client 34 | // connections, forward among the ring and replicate. 35 | type Node struct { 36 | name string 37 | logger *log.Logger 38 | serf *serf.Serf 39 | eventCh chan serf.Event 40 | quitCh chan bool 41 | ring ring 42 | replicas int 43 | nodes map[string]*peerNode 44 | data *ttlstore.DataStore 45 | replicator dataReplicator 46 | port int 47 | listener net.Listener 48 | shutdown bool 49 | server *rpcServer 50 | lock sync.Mutex 51 | config *NodeConfig 52 | gc *time.Timer 53 | } 54 | 55 | // Create creates a new Node instance and initializes all resources necessary to 56 | // be a member of the hash ring. 57 | func Create(config *NodeConfig) (*Node, error) { 58 | if err := configValidator(config); err != nil { 59 | return nil, err 60 | } 61 | 62 | node := &Node{ 63 | name: uuid.New(), 64 | eventCh: make(chan serf.Event), 65 | quitCh: make(chan bool), 66 | replicas: config.Replicas, 67 | nodes: make(map[string]*peerNode), 68 | config: config, 69 | } 70 | 71 | // RPC Endpoint 72 | node.server = newRpcServer(node) 73 | 74 | // Initialize the ring. 75 | node.ring = newConsistentRing(node.name, config.Replicas) 76 | 77 | // Set up logging. 78 | if config.LogOutput == nil { 79 | config.LogOutput = os.Stderr 80 | } 81 | node.logger = log.New(config.LogOutput, "", log.LstdFlags) 82 | 83 | // Initialize the data storage. 84 | d, err := ttlstore.New(config.MaxMemoryUsage, node.logger) 85 | if err != nil { 86 | return nil, fmt.Errorf("DataStore initialization failed: %v", err) 87 | } 88 | node.data = d 89 | 90 | // Initialize replication 91 | if config.Replicas == 1 { 92 | node.replicator = &noOpReplicator{} 93 | } else { 94 | node.replicator = newReplicator(node) 95 | } 96 | 97 | // Initialize garbage collection. 98 | node.gc = time.AfterFunc(10*time.Minute, node.garbageCollect) 99 | node.gc.Stop() 100 | 101 | // Initialize the RPC Server 102 | if err := node.initRpcServer(config.RpcPort); err != nil { 103 | return nil, fmt.Errorf("RPC server initialization failed: %v", err) 104 | } 105 | 106 | if err := node.initSerf(config.SerfConfig); err != nil { 107 | return nil, fmt.Errorf("serf initialization failed: %v", err) 108 | } 109 | 110 | return node, nil 111 | } 112 | 113 | // Start starts handling connections from other rings in the cluster as well as 114 | // clients. 115 | func (n *Node) Start() { 116 | go n.serfEventHandler() 117 | go n.listen() 118 | } 119 | 120 | // configValidator ensures that the configuration is valid and returns an error 121 | // if it is not. 122 | func configValidator(config *NodeConfig) error { 123 | if config == nil { 124 | return errors.New("must supply a non-nil config") 125 | } 126 | 127 | if config.Replicas < 1 { 128 | return errors.New("must have at least one replica") 129 | } 130 | 131 | return nil 132 | } 133 | 134 | // initSerf initializes serf and advertizes this nodes RPC port to other nodes. 135 | func (n *Node) initSerf(config *serf.Config) error { 136 | config.NodeName = n.name 137 | config.EventCh = n.eventCh 138 | 139 | // Advertize the RPC port 140 | tags := map[string]string{ 141 | externalRpcTag: strconv.Itoa(n.port), 142 | } 143 | 144 | config.Tags = tags 145 | 146 | // Start serf for discovery 147 | s, err := serf.Create(config) 148 | if err != nil { 149 | return err 150 | } 151 | 152 | n.serf = s 153 | return nil 154 | } 155 | 156 | // serfEventHandler is a dispatch method on Serf Events. Should be started in a 157 | // goroutine. 158 | func (n *Node) serfEventHandler() { 159 | for { 160 | select { 161 | case <-n.quitCh: 162 | return 163 | case e := <-n.eventCh: 164 | switch t := e.(type) { 165 | case serf.MemberEvent: 166 | switch t.EventType() { 167 | case serf.EventMemberJoin: 168 | go n.handleNodeJoins(t.Members) 169 | case serf.EventMemberLeave, serf.EventMemberFailed: 170 | go n.handleNodeLeaves(t.Members) 171 | default: 172 | // Don't handle any other event type. 173 | continue 174 | } 175 | default: 176 | n.logger.Printf("[ERROR] onecache: unexpected Serf event type %v\n", t) 177 | continue 178 | } 179 | } 180 | } 181 | } 182 | 183 | // Exit removes this Node from the ring, kills background tasks and frees any 184 | // used resources. Once this is called the Node will no longer accept 185 | // connections and can not be used. 186 | func (n *Node) Exit() error { 187 | close(n.quitCh) 188 | n.replicator.Stop() 189 | 190 | // Clean up the datastore. 191 | if err := n.data.Destroy(); err != nil { 192 | return err 193 | } 194 | 195 | // Close the listener and RPC Server. 196 | n.shutdown = true 197 | if err := n.listener.Close(); err != nil { 198 | return err 199 | } 200 | 201 | if err := n.server.shutdown(); err != nil { 202 | return err 203 | } 204 | 205 | if err := n.serf.Leave(); err != nil { 206 | return err 207 | } 208 | 209 | if err := n.serf.Shutdown(); err != nil { 210 | return err 211 | } 212 | 213 | return nil 214 | } 215 | 216 | // Join adds this Node to an existing set ring. Nodes is a list of ip:port port pairs 217 | // or DNS addresses. 218 | func (n *Node) Join(nodes []string) error { 219 | _, err := n.serf.Join(nodes, true) 220 | return err 221 | } 222 | 223 | func (n *Node) handleNodeJoins(members []serf.Member) { 224 | n.memberStateUpdates(members, false) 225 | } 226 | 227 | func (n *Node) handleNodeLeaves(members []serf.Member) { 228 | n.memberStateUpdates(members, true) 229 | } 230 | 231 | // Members should be called when serf members join or leave. Remove is set to 232 | // true if nodes left. This method updates state, including the ring to reflect 233 | // Nodes joining/leaving. 234 | func (n *Node) memberStateUpdates(members []serf.Member, remove bool) { 235 | n.lock.Lock() 236 | 237 | var peers []string 238 | for _, m := range members { 239 | p, err := buildPeerNode(m) 240 | if err != nil { 241 | n.logger.Printf("[ERROR] onecache: buildPeerNode(%v) failed: %v\n", m, err) 242 | continue 243 | } 244 | peers = append(peers, p.name) 245 | 246 | if err := n.memberStateUpdate(p, remove); err != nil { 247 | n.logger.Printf("[ERROR] onecache: memberStateUpdate(%v, %v) failed: %v\n", m, remove, err) 248 | } 249 | } 250 | n.lock.Unlock() 251 | 252 | // Want to mark after all state updates have been applied because the state of replicas can be in flux. 253 | if remove { 254 | n.replicator.NodesRemoved(peers) 255 | } else { 256 | if err := n.replicator.MarkPeersDirty(peers); err != nil { 257 | n.logger.Printf("[ERROR] onecache.replicator: %v", err) 258 | } 259 | } 260 | 261 | // The set of Nodes has changed so trigger garbage collection 262 | n.gc.Reset(10 * time.Minute) 263 | } 264 | 265 | // memberStateUpdate handles a single node joining or leaving. 266 | func (n *Node) memberStateUpdate(p *peerNode, remove bool) error { 267 | if remove { 268 | if _, exists := n.nodes[p.name]; !exists { 269 | return fmt.Errorf("delete of non-existent peer %+v", p) 270 | } 271 | 272 | n.logger.Printf("[INFO] onecache: removed peer node %v", p) 273 | delete(n.nodes, p.name) 274 | n.ring.remove(p.name) 275 | return nil 276 | } 277 | 278 | n.logger.Printf("[INFO] onecache: added peer node %v", p) 279 | n.nodes[p.name] = p 280 | n.ring.add(p.name) 281 | return nil 282 | } 283 | 284 | // buildPeerNode takes a serf member and returns a peerNode 285 | func buildPeerNode(m serf.Member) (*peerNode, error) { 286 | name := m.Name 287 | portString, ok := m.Tags[externalRpcTag] 288 | if !ok { 289 | return nil, fmt.Errorf("member joined without the rpc port set: %+v", m) 290 | } 291 | 292 | port, err := strconv.ParseInt(portString, 0, 0) 293 | if err != nil { 294 | return nil, fmt.Errorf("couldn't parse rpc port string %v for member %+v", portString, m) 295 | } 296 | 297 | return &peerNode{ 298 | name: name, 299 | addr: m.Addr.String(), 300 | port: int(port), 301 | }, nil 302 | } 303 | 304 | // garbageCollect scans for stored keys that this node doesn't own or should not 305 | // be replicated to it and removes them. 306 | func (n *Node) garbageCollect() { 307 | keys, replicating := n.getGCKeys() 308 | for _, key := range keys { 309 | n.data.Delete(key) 310 | } 311 | 312 | // Queue another round of GC since we were replicating unowned keys. 313 | if replicating { 314 | n.gc.Reset(10 * time.Minute) 315 | } 316 | } 317 | 318 | // getGCKeys returns keys that should be GC'd and whether there should be a 319 | // follow up GC because keys are being replicated. 320 | func (n *Node) getGCKeys() ([]string, bool) { 321 | keys := n.data.List() 322 | var gc []string 323 | var replicating bool 324 | for _, key := range keys { 325 | replicas := n.ring.replicaPeers(key) 326 | for _, replica := range replicas { 327 | if replica == n.name { 328 | if n.replicator.IsReplicating(key) { 329 | replicating = true 330 | continue 331 | } 332 | gc = append(gc, key) 333 | break 334 | } 335 | } 336 | } 337 | 338 | return gc, replicating 339 | } 340 | -------------------------------------------------------------------------------- /onecache/onecache_rpc.go: -------------------------------------------------------------------------------- 1 | package onecache 2 | 3 | import ( 4 | "fmt" 5 | "io" 6 | "net" 7 | "net/rpc" 8 | "sync" 9 | "time" 10 | 11 | "github.com/dadgar/onecache/ttlstore" 12 | "github.com/hashicorp/net-rpc-msgpackrpc" 13 | "github.com/hashicorp/yamux" 14 | "github.com/ugorji/go/codec" 15 | ) 16 | 17 | const ( 18 | // nodeRPCCache controls how long we keep an idle connection open to a node 19 | nodeRPCCache = 2 * time.Minute 20 | 21 | // nodeMaxStreams controsl how many idle streams we keep open to a node 22 | nodeMaxStreams = 64 23 | ) 24 | 25 | var ( 26 | mh codec.Handle = new(codec.MsgpackHandle) 27 | ) 28 | 29 | // rpcServer handles rpc connections between nodes. 30 | type rpcServer struct { 31 | n *Node 32 | l sync.Mutex 33 | versions map[string]*ReplicaVersion 34 | rpc *rpc.Server 35 | 36 | // Connection pool to other nodes 37 | connPool *connPool 38 | } 39 | 40 | // newRpcServer returns an rpc server suitable for handling RPC connections for 41 | // the passed node. 42 | func newRpcServer(n *Node) *rpcServer { 43 | s := &rpcServer{ 44 | n: n, 45 | versions: make(map[string]*ReplicaVersion), 46 | rpc: rpc.NewServer(), 47 | connPool: newPool(n.config.LogOutput, nodeRPCCache, nodeMaxStreams), 48 | } 49 | 50 | s.rpc.RegisterName("Node", s) 51 | return s 52 | } 53 | 54 | // shutdown stops the rpc server from handling new connections. 55 | func (r *rpcServer) shutdown() error { 56 | return r.connPool.Shutdown() 57 | } 58 | 59 | // handleMultiplex handles a multiplexed connection. 60 | func (r *rpcServer) handleMultiplex(conn net.Conn) { 61 | defer conn.Close() 62 | conf := yamux.DefaultConfig() 63 | conf.LogOutput = r.n.config.LogOutput 64 | server, _ := yamux.Server(conn, conf) 65 | for { 66 | sub, err := server.Accept() 67 | if err != nil { 68 | if err != io.EOF { 69 | r.n.logger.Printf("[ERR] onecache.rpc: multiplex conn accept failed: %v", err) 70 | } 71 | return 72 | } 73 | go r.handleConn(sub) 74 | } 75 | } 76 | 77 | // handleConn handles a non-multiplexed connection. 78 | func (r *rpcServer) handleConn(conn net.Conn) { 79 | rpcCodec := msgpackrpc.NewServerCodec(conn) 80 | for { 81 | select { 82 | case <-r.n.quitCh: 83 | return 84 | default: 85 | break 86 | } 87 | 88 | r.rpc.ServeRequest(rpcCodec) 89 | } 90 | } 91 | 92 | // targetNode takes the key and the explicitPeer value from the RPC and returns 93 | // the target node for the connection. 94 | func (r *rpcServer) targetNode(key, explictPeer string) string { 95 | if explictPeer != "" { 96 | return explictPeer 97 | } 98 | 99 | return r.n.ring.successor(key) 100 | } 101 | 102 | // forward forwards the rpc to the passed target node. 103 | func (r *rpcServer) forward(target, method string, args interface{}, reply interface{}) error { 104 | r.n.lock.Lock() 105 | p, ok := r.n.nodes[target] 106 | r.n.lock.Unlock() 107 | if !ok { 108 | return fmt.Errorf("could not dial unknown peer %v", target) 109 | } 110 | 111 | addr, err := net.ResolveTCPAddr("tcp", fmt.Sprintf("%s:%d", p.addr, p.port)) 112 | if err != nil { 113 | return fmt.Errorf("Failed to resolve tcp address of peer %v: %v", target, err) 114 | } 115 | 116 | err = r.connPool.RPC(addr, method, args, reply) 117 | return err 118 | } 119 | 120 | // StorageArgs is the struct used to make data related RPC calls. 121 | type StorageArgs struct { 122 | Key string 123 | Value []byte 124 | Flags int32 125 | Exp int64 126 | Cas int64 127 | 128 | // Used for increment and decrement. 129 | Delta int64 130 | 131 | // Used internally as a means to determine when to forward. 132 | ExplicitPeer string 133 | } 134 | 135 | // markClientVersion is used to mark that a key was set by a client and not by 136 | // replication. 137 | func (r *rpcServer) markClientVersion(key string) { 138 | r.l.Lock() 139 | defer r.l.Unlock() 140 | delete(r.versions, key) 141 | } 142 | 143 | // Set is used to set the value of a key. 144 | func (r *rpcServer) Set(args StorageArgs, resp *struct{}) error { 145 | target := r.targetNode(args.Key, args.ExplicitPeer) 146 | if target != r.n.name { 147 | return r.forward(target, "Node.Set", args, resp) 148 | } 149 | 150 | if err := r.n.data.Set(args.Key, args.Value, args.Exp, args.Flags); err != nil { 151 | return err 152 | } 153 | 154 | r.n.replicator.MarkDirty(args.Key) 155 | r.markClientVersion(args.Key) 156 | return nil 157 | } 158 | 159 | // Cas is used to set the value of a key using a check and set. 160 | func (r *rpcServer) Cas(args StorageArgs, resp *struct{}) error { 161 | target := r.targetNode(args.Key, args.ExplicitPeer) 162 | if target != r.n.name { 163 | return r.forward(target, "Node.Cas", args, resp) 164 | } 165 | 166 | if err := r.n.data.Cas(args.Key, args.Value, args.Exp, args.Cas, args.Flags); err != nil { 167 | return err 168 | } 169 | 170 | r.n.replicator.MarkDirty(args.Key) 171 | r.markClientVersion(args.Key) 172 | return nil 173 | } 174 | 175 | // Add adds a key only if it doesn't currently exist. 176 | func (r *rpcServer) Add(args StorageArgs, resp *struct{}) error { 177 | target := r.targetNode(args.Key, args.ExplicitPeer) 178 | if target != r.n.name { 179 | return r.forward(target, "Node.Add", args, resp) 180 | } 181 | 182 | if err := r.n.data.Add(args.Key, args.Value, args.Exp, args.Flags); err != nil { 183 | return err 184 | } 185 | 186 | r.n.replicator.MarkDirty(args.Key) 187 | r.markClientVersion(args.Key) 188 | return nil 189 | } 190 | 191 | // Replace sets the content of a key only if it exists. 192 | func (r *rpcServer) Replace(args StorageArgs, resp *struct{}) error { 193 | target := r.targetNode(args.Key, args.ExplicitPeer) 194 | if target != r.n.name { 195 | return r.forward(target, "Node.Replace", args, resp) 196 | } 197 | 198 | if err := r.n.data.Replace(args.Key, args.Value, args.Exp, args.Flags); err != nil { 199 | return err 200 | } 201 | 202 | r.n.replicator.MarkDirty(args.Key) 203 | r.markClientVersion(args.Key) 204 | return nil 205 | } 206 | 207 | // Append appends to the contents of a key. 208 | func (r *rpcServer) Append(args StorageArgs, resp *struct{}) error { 209 | target := r.targetNode(args.Key, args.ExplicitPeer) 210 | if target != r.n.name { 211 | return r.forward(target, "Node.Append", args, resp) 212 | } 213 | 214 | if err := r.n.data.Append(args.Key, args.Value); err != nil { 215 | return err 216 | } 217 | 218 | r.n.replicator.MarkDirty(args.Key) 219 | r.markClientVersion(args.Key) 220 | return nil 221 | } 222 | 223 | // Prepend prepends to the contents of a key. 224 | func (r *rpcServer) Prepend(args StorageArgs, resp *struct{}) error { 225 | target := r.targetNode(args.Key, args.ExplicitPeer) 226 | if target != r.n.name { 227 | return r.forward(target, "Node.Prepend", args, resp) 228 | } 229 | 230 | if err := r.n.data.Prepend(args.Key, args.Value); err != nil { 231 | return err 232 | } 233 | 234 | r.n.replicator.MarkDirty(args.Key) 235 | r.markClientVersion(args.Key) 236 | return nil 237 | } 238 | 239 | // Increment increments the value of a key. 240 | func (r *rpcServer) Increment(args StorageArgs, resp *int64) error { 241 | target := r.targetNode(args.Key, args.ExplicitPeer) 242 | if target != r.n.name { 243 | return r.forward(target, "Node.Increment", args, resp) 244 | } 245 | 246 | incr, err := r.n.data.Increment(args.Key, args.Delta) 247 | if err != nil { 248 | return err 249 | } 250 | *resp = incr 251 | 252 | r.n.replicator.MarkDirty(args.Key) 253 | r.markClientVersion(args.Key) 254 | return nil 255 | } 256 | 257 | // Decrement decrements the value of a key. 258 | func (r *rpcServer) Decrement(args StorageArgs, resp *int64) error { 259 | target := r.targetNode(args.Key, args.ExplicitPeer) 260 | if target != r.n.name { 261 | return r.forward(target, "Node.Decrement", args, resp) 262 | } 263 | 264 | decr, err := r.n.data.Decrement(args.Key, args.Delta) 265 | if err != nil { 266 | return err 267 | } 268 | *resp = decr 269 | 270 | r.n.replicator.MarkDirty(args.Key) 271 | r.markClientVersion(args.Key) 272 | return nil 273 | } 274 | 275 | // Contains returns whether a key exists. 276 | func (r *rpcServer) Contains(args StorageArgs, resp *bool) error { 277 | target := r.targetNode(args.Key, args.ExplicitPeer) 278 | if target != r.n.name { 279 | return r.forward(target, "Node.Contains", args, resp) 280 | } 281 | 282 | *resp = r.n.data.Contains(args.Key) 283 | r.markClientVersion(args.Key) 284 | return nil 285 | } 286 | 287 | // Touch is used to set the experation of a key. 288 | func (r *rpcServer) Touch(args StorageArgs, resp *struct{}) error { 289 | target := r.targetNode(args.Key, args.ExplicitPeer) 290 | if target != r.n.name { 291 | return r.forward(target, "Node.Touch", args, resp) 292 | } 293 | 294 | if err := r.n.data.Touch(args.Key, args.Exp); err != nil { 295 | return err 296 | } 297 | 298 | r.n.replicator.MarkTouched(args.Key) 299 | r.markClientVersion(args.Key) 300 | return nil 301 | } 302 | 303 | // Remove removes the key. 304 | func (r *rpcServer) Remove(args StorageArgs, resp *struct{}) error { 305 | target := r.targetNode(args.Key, args.ExplicitPeer) 306 | if target != r.n.name { 307 | return r.forward(target, "Node.Remove", args, resp) 308 | } 309 | 310 | if err := r.n.data.Delete(args.Key); err != nil { 311 | return err 312 | } 313 | 314 | r.n.replicator.MarkDeleted(args.Key) 315 | r.markClientVersion(args.Key) 316 | return nil 317 | } 318 | 319 | // Get returns the value of the key. 320 | func (r *rpcServer) Get(args StorageArgs, resp *ttlstore.KeyData) error { 321 | target := r.targetNode(args.Key, args.ExplicitPeer) 322 | if target != r.n.name { 323 | return r.forward(target, "Node.Get", args, resp) 324 | } 325 | 326 | data, err := r.n.data.Get(args.Key) 327 | if err != nil { 328 | return err 329 | } 330 | 331 | *resp = *data 332 | return nil 333 | } 334 | 335 | // ReplicaStorageArgs is a versioned key set by a replica peer. 336 | type ReplicaStorageArgs struct { 337 | StorageArgs 338 | ReplicaVersion 339 | } 340 | 341 | // ReplicaVersion is used to version replicated data of a key from a peer. 342 | type ReplicaVersion struct { 343 | // The peer sending the replica information 344 | Peer string 345 | 346 | // A version which should only increment 347 | Version int64 348 | } 349 | 350 | // ReplicaSet is used for setting the value of a key to a replica. It allows 351 | // setting the version of the key so that only the most updated version of the 352 | // value is stored. 353 | func (r *rpcServer) ReplicaSet(args ReplicaStorageArgs, resp *struct{}) error { 354 | target := r.targetNode(args.Key, args.ExplicitPeer) 355 | if target != r.n.name { 356 | return r.forward(target, "Node.ReplicaSet", args, resp) 357 | } 358 | 359 | r.l.Lock() 360 | defer r.l.Unlock() 361 | v := r.versions[args.Key] 362 | 363 | // Check if we are trying to set an old version. 364 | if v != nil && v.Peer == args.Peer && args.Version <= v.Version { 365 | return nil 366 | } 367 | 368 | if err := r.n.data.Set(args.Key, args.Value, args.Exp, args.Flags); err != nil { 369 | return err 370 | } 371 | 372 | // Update the replica version. 373 | r.versions[args.Key] = &args.ReplicaVersion 374 | return nil 375 | } 376 | 377 | // ReplicaTouch is used for setting the expiration of a key on a replica. It 378 | // allows setting the version of the key so that only the most updated version 379 | // of the value is stored. 380 | func (r *rpcServer) ReplicaTouch(args ReplicaStorageArgs, resp *struct{}) error { 381 | target := r.targetNode(args.Key, args.ExplicitPeer) 382 | if target != r.n.name { 383 | return r.forward(target, "Node.ReplicaTouch", args, resp) 384 | } 385 | 386 | r.l.Lock() 387 | defer r.l.Unlock() 388 | v := r.versions[args.Key] 389 | 390 | // Check if we are trying to set an old version. 391 | if v != nil && v.Peer == args.Peer && args.Version <= v.Version { 392 | return nil 393 | } 394 | 395 | if err := r.n.data.Touch(args.Key, args.Exp); err != nil { 396 | return err 397 | } 398 | 399 | // Update the replica version. 400 | r.versions[args.Key] = &args.ReplicaVersion 401 | return nil 402 | } 403 | -------------------------------------------------------------------------------- /onecache/onecache_rpc_test.go: -------------------------------------------------------------------------------- 1 | package onecache 2 | 3 | import ( 4 | "reflect" 5 | "testing" 6 | "time" 7 | 8 | "github.com/dadgar/onecache/ttlstore" 9 | ) 10 | 11 | func TestStoreRemote(t *testing.T) { 12 | nodes, err := connectedTestNodes(2, 1) 13 | defer tearDownNodes(nodes) 14 | if err != nil { 15 | t.Fatalf("connectedTestNodes(2, 1) failed: %v", err) 16 | } 17 | 18 | n1 := nodes[0] 19 | n2 := nodes[1] 20 | 21 | expV := []byte("bar") 22 | args := StorageArgs{ 23 | Key: "foo", 24 | Value: expV, 25 | Flags: 0, 26 | Exp: 0, 27 | ExplicitPeer: n2.name, 28 | } 29 | 30 | var resp struct{} 31 | if err := n1.server.Set(args, &resp); err != nil { 32 | t.Errorf("Set(%+v) returned an error: %v", args, err) 33 | } 34 | 35 | // Check the key is actually stored. 36 | data, err := n2.data.Get(args.Key) 37 | if err != nil { 38 | t.Errorf("n2.data.Get(%v) returned error: %v", args.Key, err) 39 | t.FailNow() 40 | } 41 | 42 | if !reflect.DeepEqual(data.Data, expV) { 43 | t.Errorf("Stored incorrect value: want %v; got %v", expV, data.Data) 44 | } 45 | } 46 | 47 | func TestCasRemoteInvalidRemote(t *testing.T) { 48 | nodes, err := connectedTestNodes(2, 1) 49 | defer tearDownNodes(nodes) 50 | if err != nil { 51 | t.Fatalf("connectedTestNodes(2, 1) failed: %v", err) 52 | } 53 | 54 | n1 := nodes[0] 55 | n2 := nodes[1] 56 | 57 | expV := []byte("bar") 58 | args := StorageArgs{ 59 | Key: "foo", 60 | Value: expV, 61 | Cas: 10, 62 | Flags: 0, 63 | Exp: 0, 64 | ExplicitPeer: n2.name, 65 | } 66 | 67 | var resp struct{} 68 | if err := n1.server.Cas(args, &resp); err == nil { 69 | t.Errorf("Cas(%+v) returned %v; expected error", args) 70 | } 71 | } 72 | 73 | func TestContainsRemote(t *testing.T) { 74 | nodes, err := connectedTestNodes(2, 1) 75 | defer tearDownNodes(nodes) 76 | if err != nil { 77 | t.Fatalf("connectedTestNodes(2, 1) failed: %v", err) 78 | } 79 | 80 | n1 := nodes[0] 81 | n2 := nodes[1] 82 | 83 | // Store a value on the second node 84 | k := "foo" 85 | if err := n2.data.Set(k, []byte("bar"), 0, 0); err != nil { 86 | t.Fatalf("Set(%v) failed: %v", k, err) 87 | } 88 | 89 | var resp bool 90 | args := StorageArgs{Key: k, ExplicitPeer: n2.name} 91 | if err := n1.server.Contains(args, &resp); err != nil { 92 | t.Errorf("Contains(%+v) returned an error %v", args, err) 93 | } else if !resp { 94 | t.Errorf("Contains(%+v) returned false; want true (%v)", args, resp) 95 | } 96 | } 97 | 98 | func TestContainsRemoteInvalid(t *testing.T) { 99 | nodes, err := connectedTestNodes(2, 1) 100 | defer tearDownNodes(nodes) 101 | if err != nil { 102 | t.Fatalf("connectedTestNodes(2, 1) failed: %v", err) 103 | } 104 | 105 | n1 := nodes[0] 106 | n2 := nodes[1] 107 | 108 | var resp bool 109 | args := StorageArgs{Key: "foo", ExplicitPeer: n2.name} 110 | if err := n1.server.Contains(args, &resp); err != nil { 111 | t.Fatalf("Contains(%+v) returned an error: %v", args, err) 112 | } 113 | 114 | if resp { 115 | t.Errorf("Contains(%+v) returned %v; want false", args, resp) 116 | } 117 | } 118 | 119 | func TestTouchRemote(t *testing.T) { 120 | nodes, err := connectedTestNodes(2, 1) 121 | defer tearDownNodes(nodes) 122 | if err != nil { 123 | t.Fatalf("connectedTestNodes(2, 1) failed: %v", err) 124 | } 125 | 126 | n1 := nodes[0] 127 | n2 := nodes[1] 128 | 129 | // Store a value on the second node 130 | k := "foo" 131 | if err := n2.data.Set(k, []byte("bar"), 0, 0); err != nil { 132 | t.Fatalf("Set(%v) failed: %v", k, err) 133 | } 134 | 135 | var resp struct{} 136 | var exp int64 = time.Now().Add(1 * time.Hour).Unix() 137 | args := StorageArgs{Key: k, Exp: exp, ExplicitPeer: n2.name} 138 | if err := n1.server.Touch(args, &resp); err != nil { 139 | t.Errorf("Touch(%+v) returned an error %v", args, err) 140 | } 141 | 142 | // Check the new experation is correct. 143 | data, err := n2.data.Get(k) 144 | if err != nil { 145 | t.Errorf("Get(%v) failed: %v", k, err) 146 | t.FailNow() 147 | } 148 | 149 | if data.Exp != exp { 150 | t.Errorf("Experation is %v; want %v", data.Exp, exp) 151 | } 152 | } 153 | 154 | func TestTouchRemoteInvalid(t *testing.T) { 155 | nodes, err := connectedTestNodes(2, 1) 156 | defer tearDownNodes(nodes) 157 | if err != nil { 158 | t.Fatalf("connectedTestNodes(2, 1) failed: %v", err) 159 | } 160 | 161 | n1 := nodes[0] 162 | n2 := nodes[1] 163 | 164 | var resp struct{} 165 | args := StorageArgs{Key: "foo", Exp: 10, ExplicitPeer: n2.name} 166 | if err := n1.server.Touch(args, &resp); err == nil { 167 | t.Errorf("Touch(%+v) didn't return error", args) 168 | } 169 | } 170 | 171 | func TestRemoveRemote(t *testing.T) { 172 | nodes, err := connectedTestNodes(2, 1) 173 | defer tearDownNodes(nodes) 174 | if err != nil { 175 | t.Fatalf("connectedTestNodes(2, 1) failed: %v", err) 176 | } 177 | 178 | n1 := nodes[0] 179 | n2 := nodes[1] 180 | 181 | // Store a value on the second node 182 | k := "foo" 183 | if err := n2.data.Set(k, []byte("bar"), 0, 0); err != nil { 184 | t.Fatalf("Set(%v) failed: %v", k, err) 185 | } 186 | 187 | var resp struct{} 188 | args := StorageArgs{Key: k, ExplicitPeer: n2.name} 189 | if err := n1.server.Remove(args, &resp); err != nil { 190 | t.Errorf("Remove(%+v) returned an error %v", args, err) 191 | } 192 | 193 | // Check the key is actually deleted. 194 | if n2.data.Contains(k) { 195 | t.Errorf("Contains(%v) returned true; want false", k) 196 | } 197 | } 198 | 199 | func TestRemoveRemoteInvalid(t *testing.T) { 200 | nodes, err := connectedTestNodes(2, 1) 201 | defer tearDownNodes(nodes) 202 | if err != nil { 203 | t.Fatalf("connectedTestNodes(2, 1) failed: %v", err) 204 | } 205 | 206 | n1 := nodes[0] 207 | n2 := nodes[1] 208 | 209 | var resp struct{} 210 | args := StorageArgs{Key: "foo", ExplicitPeer: n2.name} 211 | if err := n1.server.Remove(args, &resp); err == nil { 212 | t.Errorf("Remove(%+v) didn't return error", args) 213 | } 214 | } 215 | 216 | func TestGetRemote(t *testing.T) { 217 | nodes, err := connectedTestNodes(2, 1) 218 | defer tearDownNodes(nodes) 219 | if err != nil { 220 | t.Fatalf("connectedTestNodes(2, 1) failed: %v", err) 221 | } 222 | 223 | n1 := nodes[0] 224 | n2 := nodes[1] 225 | 226 | // Store a value on the second node 227 | k := "foo" 228 | expV := []byte("bar") 229 | if err := n2.data.Set(k, expV, 0, 0); err != nil { 230 | t.Fatalf("Set(%v) failed: %v", k, err) 231 | } 232 | 233 | var resp ttlstore.KeyData 234 | args := StorageArgs{Key: k, ExplicitPeer: n2.name} 235 | if n1.server.Get(args, &resp); err != nil { 236 | t.Errorf("Get(%+v) returned an error %v", args, err) 237 | t.FailNow() 238 | } 239 | 240 | // Check the value returned is correct. 241 | if !reflect.DeepEqual(resp.Data, expV) { 242 | t.Errorf("Get(%v) returned %v; want %v", args, resp.Data, expV) 243 | } 244 | } 245 | 246 | func TestGetRemoteInvalid(t *testing.T) { 247 | nodes, err := connectedTestNodes(2, 1) 248 | defer tearDownNodes(nodes) 249 | if err != nil { 250 | t.Fatalf("connectedTestNodes(2, 1) failed: %v", err) 251 | } 252 | 253 | n1 := nodes[0] 254 | n2 := nodes[1] 255 | 256 | var resp ttlstore.KeyData 257 | args := StorageArgs{Key: "foo", ExplicitPeer: n2.name} 258 | if err := n1.server.Get(args, &resp); err == nil { 259 | t.Errorf("Get(%+v) didn't return error", args) 260 | } 261 | } 262 | -------------------------------------------------------------------------------- /onecache/onecache_test.go: -------------------------------------------------------------------------------- 1 | package onecache 2 | 3 | import ( 4 | "fmt" 5 | "reflect" 6 | "testing" 7 | "time" 8 | ) 9 | 10 | func TestCreateInvalid(t *testing.T) { 11 | if _, err := Create(nil); err == nil { 12 | t.Error("Create(nil) should have failed") 13 | } 14 | } 15 | 16 | func TestNewRpcServerAnyPort(t *testing.T) { 17 | n := &Node{} 18 | 19 | // A zero value port allows any port binding. 20 | err := n.initRpcServer(0) 21 | 22 | if err != nil { 23 | t.Error(err) 24 | } 25 | 26 | if n.listener == nil { 27 | t.Error("newRpcServer(0) did not create a listener") 28 | } 29 | 30 | if n.port <= 0 { 31 | t.Error("newRpcServer(0) did not bind to an available port") 32 | } 33 | } 34 | 35 | func TestStartRpcServerSpecificPort(t *testing.T) { 36 | n := &Node{} 37 | 38 | // A zero value port allows any port binding. 39 | port := 9100 40 | err := n.initRpcServer(port) 41 | 42 | if err != nil { 43 | t.Error(err) 44 | } 45 | 46 | if n.listener == nil { 47 | t.Errorf("newRpcServer(%v) did not create a listener", port) 48 | } 49 | 50 | if n.port != port { 51 | t.Errorf("newRpcServer(%v) bound to %v; want %v", port, n.port, port) 52 | } 53 | } 54 | 55 | func TestStartRpcServerNegativePort(t *testing.T) { 56 | n := &Node{} 57 | 58 | // A zero value port allows any port binding. 59 | if err := n.initRpcServer(-1); err == nil { 60 | t.Error("newRpcServer(-1) should return an error") 61 | } 62 | } 63 | 64 | func TestExit(t *testing.T) { 65 | n, err := defaultTestNode() 66 | if err != nil { 67 | t.Fatalf("defaultTestNode() failed: %v", err) 68 | } 69 | 70 | if err := n.Exit(); err != nil { 71 | t.Errorf("Exit() returned an error: %v", err) 72 | } 73 | } 74 | 75 | func TestJoinInvalid(t *testing.T) { 76 | n, err := defaultTestNode() 77 | if err != nil { 78 | t.Fatalf("defaultTestNode() failed: %v", err) 79 | } 80 | 81 | others := []string{"junk:7337"} 82 | if err := n.Join(others); err == nil { 83 | t.Errorf("Join(%v) should have failed", others) 84 | } 85 | } 86 | 87 | func TestJoinValid(t *testing.T) { 88 | // Create two Nodes 89 | config := buildTestConfig(2, 0, 0) 90 | n, err := Create(config) 91 | if err != nil { 92 | t.Fatalf("Create(%+v) failed: %v", config, err) 93 | } 94 | 95 | config2 := buildTestConfig(2, 0, 0) 96 | n2, err := Create(config2) 97 | if err != nil { 98 | t.Fatalf("Create(%+v) failed: %v", config, err) 99 | } 100 | 101 | time.Sleep(10 * time.Millisecond) 102 | 103 | if len(n.serf.Members()) != 1 { 104 | t.Fatalf("s1 members: %d", len(n.serf.Members())) 105 | } 106 | 107 | if len(n2.serf.Members()) != 1 { 108 | t.Fatalf("s2 members: %d", len(n2.serf.Members())) 109 | } 110 | 111 | others := []string{fmt.Sprintf("127.0.0.1:%d", n2.serf.LocalMember().Port)} 112 | if err := n.Join(others); err != nil { 113 | t.Errorf("Join(%v) failed: %v", others, err) 114 | } 115 | 116 | time.Sleep(10 * time.Millisecond) 117 | 118 | if len(n.serf.Members()) != 2 { 119 | t.Errorf("s1 members: %d", len(n.serf.Members())) 120 | } 121 | 122 | if len(n2.serf.Members()) != 2 { 123 | t.Errorf("s2 members: %d", len(n2.serf.Members())) 124 | } 125 | } 126 | 127 | func TestBuildPeerNameValid(t *testing.T) { 128 | n, err := defaultTestNode() 129 | if err != nil { 130 | t.Fatalf("defaultTestNode() failed: %v", err) 131 | } 132 | 133 | if len(n.serf.Members()) != 1 { 134 | t.Fatalf("serf members: %d", len(n.serf.Members())) 135 | } 136 | 137 | // Get the Serf member. 138 | m := n.serf.Members()[0] 139 | 140 | expected := peerNode{ 141 | name: n.name, 142 | addr: m.Addr.String(), 143 | port: n.port, 144 | } 145 | 146 | peer, err := buildPeerNode(m) 147 | if err != nil { 148 | t.Errorf("buildPeerNode(%v) failed: %v", m, err) 149 | } 150 | 151 | if !reflect.DeepEqual(expected, *peer) { 152 | t.Errorf("buildPeerNode(%v) returned %v; want %v", m, *peer, expected) 153 | } 154 | } 155 | 156 | func TestNode_GarbageCollection(t *testing.T) { 157 | n, err := defaultTestNode() 158 | if err != nil { 159 | t.Fatalf("defaultTestNode() failed: %v", err) 160 | } 161 | 162 | // Inject a key that should be GC'd, shouldn't be and those that are 163 | // replicating. 164 | keys := []string{"foo", "bar", "baz"} 165 | for _, k := range keys { 166 | if err := n.data.Set(k, []byte(k), 0, 0); err != nil { 167 | t.Fatalf("Set() failed: %v", err) 168 | } 169 | } 170 | 171 | // Mock the replicator such that the first key is being replicated. 172 | n.replicator = newMockReplicator(func(key string) bool { 173 | if key == keys[0] { 174 | return true 175 | } 176 | return false 177 | }) 178 | 179 | // Mock the ring such that the second key shouldn't be replicated. 180 | n.ring = &mockRing{ 181 | replicaPeersImpl: func(key string) []string { 182 | if key == keys[1] { 183 | return nil 184 | } 185 | 186 | return []string{n.name} 187 | }, 188 | } 189 | 190 | gc, replicating := n.getGCKeys() 191 | if !replicating { 192 | t.Fatal("getGCKeys() should indicate that a follow up gc is needed") 193 | } 194 | 195 | if len(gc) != 1 || gc[0] != keys[2] { 196 | t.Fatalf("getGCKeys() returned wrong key") 197 | } 198 | } 199 | -------------------------------------------------------------------------------- /onecache/pool.go: -------------------------------------------------------------------------------- 1 | package onecache 2 | 3 | import ( 4 | "container/list" 5 | "fmt" 6 | "io" 7 | "net" 8 | "net/rpc" 9 | "sync" 10 | "sync/atomic" 11 | "time" 12 | 13 | "github.com/hashicorp/net-rpc-msgpackrpc" 14 | "github.com/hashicorp/yamux" 15 | ) 16 | 17 | // streamClient is used to wrap a stream with an RPC client 18 | type streamClient struct { 19 | stream net.Conn 20 | codec rpc.ClientCodec 21 | } 22 | 23 | func (sc *streamClient) Close() { 24 | sc.stream.Close() 25 | } 26 | 27 | // conn is a pooled connection to a Nomad server 28 | type conn struct { 29 | refCount int32 30 | shouldClose int32 31 | 32 | addr net.Addr 33 | session *yamux.Session 34 | lastUsed time.Time 35 | 36 | pool *connPool 37 | 38 | clients *list.List 39 | clientLock sync.Mutex 40 | } 41 | 42 | // markForUse does all the bookkeeping required to ready a connection for use. 43 | func (c *conn) markForUse() { 44 | c.lastUsed = time.Now() 45 | atomic.AddInt32(&c.refCount, 1) 46 | } 47 | 48 | func (c *conn) Close() error { 49 | return c.session.Close() 50 | } 51 | 52 | // getClient is used to get a cached or new client 53 | func (c *conn) getClient() (*streamClient, error) { 54 | // Check for cached client 55 | c.clientLock.Lock() 56 | front := c.clients.Front() 57 | if front != nil { 58 | c.clients.Remove(front) 59 | } 60 | c.clientLock.Unlock() 61 | if front != nil { 62 | return front.Value.(*streamClient), nil 63 | } 64 | 65 | // Open a new session 66 | stream, err := c.session.Open() 67 | if err != nil { 68 | return nil, err 69 | } 70 | 71 | // Create a client codec 72 | codec := msgpackrpc.NewCodec(false, false, stream) 73 | 74 | // Return a new stream client 75 | sc := &streamClient{ 76 | stream: stream, 77 | codec: codec, 78 | } 79 | return sc, nil 80 | } 81 | 82 | // returnStream is used when done with a stream 83 | // to allow re-use by a future RPC 84 | func (c *conn) returnClient(client *streamClient) { 85 | didSave := false 86 | c.clientLock.Lock() 87 | if c.clients.Len() < c.pool.maxStreams && atomic.LoadInt32(&c.shouldClose) == 0 { 88 | c.clients.PushFront(client) 89 | didSave = true 90 | } 91 | c.clientLock.Unlock() 92 | if !didSave { 93 | client.Close() 94 | } 95 | } 96 | 97 | // connPool is used to maintain a connection pool to other 98 | // Nomad servers. This is used to reduce the latency of 99 | // RPC requests between servers. It is only used to pool 100 | // connections in the rpcNomad mode. Raft connections 101 | // are pooled separately. 102 | type connPool struct { 103 | sync.Mutex 104 | 105 | // LogOutput is used to control logging 106 | logOutput io.Writer 107 | 108 | // The maximum time to keep a connection open 109 | maxTime time.Duration 110 | 111 | // The maximum number of open streams to keep 112 | maxStreams int 113 | 114 | // Pool maps an address to a open connection 115 | pool map[string]*conn 116 | 117 | // limiter is used to throttle the number of connect attempts 118 | // to a given address. The first thread will attempt a connection 119 | // and put a channel in here, which all other threads will wait 120 | // on to close. 121 | limiter map[string]chan struct{} 122 | 123 | // Used to indicate the pool is shutdown 124 | shutdown bool 125 | shutdownCh chan struct{} 126 | } 127 | 128 | // newPool is used to make a new connection pool 129 | // Maintain at most one connection per host, for up to maxTime. 130 | // Set maxTime to 0 to disable reaping. maxStreams is used to control 131 | // the number of idle streams allowed. 132 | func newPool(logOutput io.Writer, maxTime time.Duration, maxStreams int) *connPool { 133 | pool := &connPool{ 134 | logOutput: logOutput, 135 | maxTime: maxTime, 136 | maxStreams: maxStreams, 137 | pool: make(map[string]*conn), 138 | limiter: make(map[string]chan struct{}), 139 | shutdownCh: make(chan struct{}), 140 | } 141 | if maxTime > 0 { 142 | go pool.reap() 143 | } 144 | return pool 145 | } 146 | 147 | // Shutdown is used to close the connection pool 148 | func (p *connPool) Shutdown() error { 149 | p.Lock() 150 | defer p.Unlock() 151 | 152 | for _, conn := range p.pool { 153 | conn.Close() 154 | } 155 | p.pool = make(map[string]*conn) 156 | 157 | if p.shutdown { 158 | return nil 159 | } 160 | p.shutdown = true 161 | close(p.shutdownCh) 162 | return nil 163 | } 164 | 165 | // Acquire is used to get a connection that is pooled or to return a new 166 | // connection 167 | func (p *connPool) acquire(addr net.Addr) (*conn, error) { 168 | // Check to see if there's a pooled connection available. This is up 169 | // here since it should the the vastly more common case than the rest 170 | // of the code here. 171 | p.Lock() 172 | s := addr.String() 173 | c := p.pool[s] 174 | if c != nil { 175 | c.markForUse() 176 | p.Unlock() 177 | return c, nil 178 | } 179 | 180 | // If not (while we are still locked), set up the throttling structure 181 | // for this address, which will make everyone else wait until our 182 | // attempt is done. 183 | var wait chan struct{} 184 | var ok bool 185 | if wait, ok = p.limiter[addr.String()]; !ok { 186 | wait = make(chan struct{}) 187 | p.limiter[addr.String()] = wait 188 | } 189 | isLeadThread := !ok 190 | p.Unlock() 191 | 192 | // If we are the lead thread, make the new connection and then wake 193 | // everybody else up to see if we got it. 194 | if isLeadThread { 195 | c, err := p.getNewConn(addr) 196 | p.Lock() 197 | delete(p.limiter, addr.String()) 198 | close(wait) 199 | if err != nil { 200 | p.Unlock() 201 | return nil, err 202 | } 203 | 204 | p.pool[addr.String()] = c 205 | p.Unlock() 206 | return c, nil 207 | } 208 | 209 | // Otherwise, wait for the lead thread to attempt the connection 210 | // and use what's in the pool at that point. 211 | select { 212 | case <-p.shutdownCh: 213 | return nil, fmt.Errorf("rpc error: shutdown") 214 | case <-wait: 215 | } 216 | 217 | // See if the lead thread was able to get us a connection. 218 | p.Lock() 219 | if c := p.pool[addr.String()]; c != nil { 220 | c.markForUse() 221 | p.Unlock() 222 | return c, nil 223 | } 224 | 225 | p.Unlock() 226 | return nil, fmt.Errorf("rpc error: lead thread didn't get connection") 227 | } 228 | 229 | // getNewConn is used to return a new connection 230 | func (p *connPool) getNewConn(addr net.Addr) (*conn, error) { 231 | // Try to dial the conn 232 | con, err := net.DialTimeout("tcp", addr.String(), 10*time.Second) 233 | if err != nil { 234 | return nil, err 235 | } 236 | 237 | // Cast to TCPConn 238 | if tcp, ok := con.(*net.TCPConn); ok { 239 | tcp.SetKeepAlive(true) 240 | tcp.SetNoDelay(true) 241 | } 242 | 243 | // Write the multiplex byte to set the mode 244 | if _, err := con.Write([]byte{byte(rpcInternal)}); err != nil { 245 | con.Close() 246 | return nil, err 247 | } 248 | 249 | // Setup the logger 250 | conf := yamux.DefaultConfig() 251 | conf.LogOutput = p.logOutput 252 | 253 | // Create a multiplexed session 254 | session, err := yamux.Client(con, conf) 255 | if err != nil { 256 | con.Close() 257 | return nil, err 258 | } 259 | 260 | // Wrap the connection 261 | c := &conn{ 262 | refCount: 1, 263 | addr: addr, 264 | session: session, 265 | clients: list.New(), 266 | lastUsed: time.Now(), 267 | pool: p, 268 | } 269 | return c, nil 270 | } 271 | 272 | // clearConn is used to clear any cached connection, potentially in response to an erro 273 | func (p *connPool) clearConn(conn *conn) { 274 | // Ensure returned streams are closed 275 | atomic.StoreInt32(&conn.shouldClose, 1) 276 | 277 | // Clear from the cache 278 | p.Lock() 279 | if c, ok := p.pool[conn.addr.String()]; ok && c == conn { 280 | delete(p.pool, conn.addr.String()) 281 | } 282 | p.Unlock() 283 | 284 | // Close down immediately if idle 285 | if refCount := atomic.LoadInt32(&conn.refCount); refCount == 0 { 286 | conn.Close() 287 | } 288 | } 289 | 290 | // releaseConn is invoked when we are done with a conn to reduce the ref count 291 | func (p *connPool) releaseConn(conn *conn) { 292 | refCount := atomic.AddInt32(&conn.refCount, -1) 293 | if refCount == 0 && atomic.LoadInt32(&conn.shouldClose) == 1 { 294 | conn.Close() 295 | } 296 | } 297 | 298 | // getClient is used to get a usable client for an address. 299 | func (p *connPool) getClient(addr net.Addr) (*conn, *streamClient, error) { 300 | retries := 0 301 | START: 302 | // Try to get a conn first 303 | conn, err := p.acquire(addr) 304 | if err != nil { 305 | return nil, nil, fmt.Errorf("failed to get conn: %v", err) 306 | } 307 | 308 | // Get a client 309 | client, err := conn.getClient() 310 | if err != nil { 311 | p.clearConn(conn) 312 | p.releaseConn(conn) 313 | 314 | // Try to redial, possible that the TCP session closed due to timeout 315 | if retries == 0 { 316 | retries++ 317 | goto START 318 | } 319 | return nil, nil, fmt.Errorf("failed to start stream: %v", err) 320 | } 321 | return conn, client, nil 322 | } 323 | 324 | // RPC is used to make an RPC call to a remote host 325 | func (p *connPool) RPC(addr net.Addr, method string, args interface{}, reply interface{}) error { 326 | // Get a usable client 327 | conn, sc, err := p.getClient(addr) 328 | if err != nil { 329 | return fmt.Errorf("rpc error: %v", err) 330 | } 331 | 332 | // Make the RPC call 333 | err = msgpackrpc.CallWithCodec(sc.codec, method, args, reply) 334 | if err != nil { 335 | sc.Close() 336 | p.releaseConn(conn) 337 | return fmt.Errorf("rpc error: %v", err) 338 | } 339 | 340 | // Done with the connection 341 | conn.returnClient(sc) 342 | p.releaseConn(conn) 343 | return nil 344 | } 345 | 346 | // Reap is used to close conns open over maxTime 347 | func (p *connPool) reap() { 348 | for { 349 | // Sleep for a while 350 | select { 351 | case <-p.shutdownCh: 352 | return 353 | case <-time.After(time.Second): 354 | } 355 | 356 | // Reap all old conns 357 | p.Lock() 358 | var removed []string 359 | now := time.Now() 360 | for host, conn := range p.pool { 361 | // Skip recently used connections 362 | if now.Sub(conn.lastUsed) < p.maxTime { 363 | continue 364 | } 365 | 366 | // Skip connections with active streams 367 | if atomic.LoadInt32(&conn.refCount) > 0 { 368 | continue 369 | } 370 | 371 | // Close the conn 372 | conn.Close() 373 | 374 | // Remove from pool 375 | removed = append(removed, host) 376 | } 377 | for _, host := range removed { 378 | delete(p.pool, host) 379 | } 380 | p.Unlock() 381 | } 382 | } 383 | -------------------------------------------------------------------------------- /onecache/replicator.go: -------------------------------------------------------------------------------- 1 | package onecache 2 | 3 | import ( 4 | "errors" 5 | "fmt" 6 | "sync" 7 | "time" 8 | ) 9 | 10 | type dataReplicator interface { 11 | // Stop stops replication. 12 | Stop() 13 | 14 | // NodesRemoved should be called when a node is removed from the ring to stop 15 | // replication. 16 | NodesRemoved(nodes []string) 17 | 18 | // Marks a key's value and metadata as dirty, causing it to be replicated. 19 | MarkDirty(key string) 20 | 21 | // Marks a key as being touched. Causes the new experation time to be replicated. 22 | MarkTouched(key string) 23 | 24 | // Marks a key as deleted and replicates the deletion. 25 | MarkDeleted(key string) 26 | 27 | // Marks a peer as dirty. Causes all keys to be replicated to the node. An error is returned if the peer doesn't 28 | // exist. 29 | MarkPeersDirty(peers []string) error 30 | 31 | // IsReplicating returns whether the passed key is being replicated. 32 | IsReplicating(key string) bool 33 | } 34 | 35 | // noOpReplicator is to be used when the node is running singularly and is not part of a cluster. 36 | type noOpReplicator struct{} 37 | 38 | func (nop *noOpReplicator) Stop() {} 39 | func (nop *noOpReplicator) IsReplicating(key string) bool { return false } 40 | func (nop *noOpReplicator) MarkDirty(key string) {} 41 | func (nop *noOpReplicator) MarkTouched(key string) {} 42 | func (nop *noOpReplicator) MarkDeleted(key string) {} 43 | func (nop *noOpReplicator) MarkPeersDirty(peers []string) error { return nil } 44 | func (nop *noOpReplicator) NodesRemoved(nodes []string) {} 45 | 46 | // replicator replicates data to peers as they are dirtied. 47 | type replicator struct { 48 | n *Node 49 | peers map[string]*peerSync 50 | quit chan struct{} 51 | lock sync.Mutex 52 | } 53 | 54 | // replicationTask identifies what replication is needed. 55 | type replicationTask byte 56 | 57 | const ( 58 | task_DELETE replicationTask = iota // Key deleted 59 | task_DIRTY_VALUE = iota // Key value changed 60 | task_TOUCH = iota // Key touched 61 | task_BACKFILL = iota // Key has new leader 62 | ) 63 | 64 | // NewReplicator returns a replicator for the passed node. 65 | func newReplicator(n *Node) *replicator { 66 | return &replicator{ 67 | n: n, 68 | peers: make(map[string]*peerSync), 69 | quit: make(chan struct{}), 70 | } 71 | } 72 | 73 | // Stop stops replication. 74 | func (r *replicator) Stop() { 75 | close(r.quit) 76 | } 77 | 78 | // IsReplicating returns if the passed key is being replicated. 79 | func (r *replicator) IsReplicating(key string) bool { 80 | r.lock.Lock() 81 | defer r.lock.Unlock() 82 | 83 | for _, peer := range r.peers { 84 | if peer.isReplicating(key) { 85 | return true 86 | } 87 | } 88 | 89 | return false 90 | } 91 | 92 | // NodesRemoved removes the nodes from being replicated to. 93 | func (r *replicator) NodesRemoved(nodes []string) { 94 | r.lock.Lock() 95 | defer r.lock.Unlock() 96 | 97 | for _, node := range nodes { 98 | peer, ok := r.peers[node] 99 | if !ok { 100 | return 101 | } 102 | 103 | peer.stop() 104 | delete(r.peers, node) 105 | } 106 | } 107 | 108 | // Marks a key's value and metadata as dirty for eventually replication to peers. 109 | func (r *replicator) MarkDirty(key string) { 110 | for _, p := range r.n.ring.replicaPeers(key) { 111 | if p == r.n.name { 112 | continue 113 | } 114 | r.getPeerSync(p).markKey(key, task_DIRTY_VALUE) 115 | } 116 | } 117 | 118 | // Marks a key's experation as dirty for eventually replication to peers. 119 | func (r *replicator) MarkTouched(key string) { 120 | for _, p := range r.n.ring.replicaPeers(key) { 121 | if p == r.n.name { 122 | continue 123 | } 124 | r.getPeerSync(p).markKey(key, task_TOUCH) 125 | } 126 | } 127 | 128 | // Marks a key as deleted for eventually replication to peers. 129 | func (r *replicator) MarkDeleted(key string) { 130 | for _, p := range r.n.ring.replicaPeers(key) { 131 | if p == r.n.name { 132 | continue 133 | } 134 | r.getPeerSync(p).markKey(key, task_DELETE) 135 | } 136 | } 137 | 138 | // Marks peers as dirty causing all keys owned by this node to be replicated to 139 | // it if it is a replica, otherwise it is a no-op. Returns an error if a peer is 140 | // unknown. 141 | func (r *replicator) MarkPeersDirty(peers []string) error { 142 | // Determine what keys where owned by this node before the peers were added. 143 | oldRing := r.n.ring.ringCopy() 144 | for _, p := range peers { 145 | oldRing.remove(p) 146 | } 147 | owned := oldRing.getOwnedKeys(r.n.data.List()) 148 | 149 | // Fast path 150 | if len(owned) == 0 { 151 | return nil 152 | } 153 | 154 | for _, peer := range peers { 155 | if peer == r.n.name { 156 | continue 157 | } 158 | 159 | r.n.lock.Lock() 160 | if _, ok := r.n.nodes[peer]; !ok { 161 | r.n.lock.Unlock() 162 | return fmt.Errorf("peer %v not found in ring", peer) 163 | } 164 | r.n.lock.Unlock() 165 | 166 | // Replicate all keys if peer is part of this nodes replica peers. 167 | if r.isReplicaPeer(peer) { 168 | r.getPeerSync(peer).markKeys(owned, task_DIRTY_VALUE) 169 | } 170 | 171 | // If the peer is taking ownership of some of our keys, we send the set 172 | // of keys it owns and mark those for replication. 173 | if taken := r.takesOwnership(peer); len(taken) != 0 { 174 | r.getPeerSync(peer).markKeys(taken, task_BACKFILL) 175 | } 176 | } 177 | 178 | return nil 179 | } 180 | 181 | // isReplicaPeer returns whether the passed peer is part of the replica set. 182 | func (r *replicator) isReplicaPeer(peer string) bool { 183 | for _, r := range r.n.ring.replicaPeers(r.n.name) { 184 | if r == peer { 185 | return true 186 | } 187 | } 188 | 189 | return false 190 | } 191 | 192 | // takesOwnership returns the keys in which the passed peer is now the owner 193 | // of. If the returned list is empty, the peer didn't take ownership of any keys. 194 | func (r *replicator) takesOwnership(peer string) []string { 195 | // Check that we are its successor. 196 | if r.n.ring.successor(peer) != r.n.name { 197 | return nil 198 | } 199 | 200 | return r.n.ring.ownedKeysByNode(r.n.data.List(), peer) 201 | } 202 | 203 | // getPeerSync returns the peer if it has been creaated, and other wise creates 204 | // it. 205 | func (r *replicator) getPeerSync(peer string) *peerSync { 206 | r.lock.Lock() 207 | defer r.lock.Unlock() 208 | p, ok := r.peers[peer] 209 | if ok { 210 | return p 211 | } 212 | 213 | p = newPeerSync(peer, r) 214 | r.peers[peer] = p 215 | return p 216 | } 217 | 218 | // peerSync is used to replicate keys to a particular peer 219 | type peerSync struct { 220 | name string 221 | replicating string 222 | tasks map[string][]replicationTask // A map of dirty keys to an operations log. 223 | queue chan string 224 | endSync chan struct{} 225 | r *replicator 226 | lock sync.Mutex 227 | } 228 | 229 | // newPeerSync returns a peerSync for replicating to the passed peer. 230 | func newPeerSync(peer string, r *replicator) *peerSync { 231 | p := &peerSync{ 232 | name: peer, 233 | tasks: make(map[string][]replicationTask), 234 | queue: make(chan string, 100), 235 | endSync: make(chan struct{}), 236 | r: r, 237 | } 238 | go p.replicate() 239 | return p 240 | } 241 | 242 | // Stop stops replication to the peer. 243 | func (p *peerSync) stop() { 244 | close(p.endSync) 245 | } 246 | 247 | // isReplicating returns if the key is being replicated to the passed peer. 248 | func (p *peerSync) isReplicating(key string) bool { 249 | p.lock.Lock() 250 | defer p.lock.Unlock() 251 | _, ok := p.tasks[key] 252 | return ok || p.replicating == key 253 | } 254 | 255 | // markKey marks the key for replication. 256 | func (p *peerSync) markKey(key string, status replicationTask) { 257 | p.lock.Lock() 258 | p.tasks[key] = append(p.tasks[key], status) 259 | p.lock.Unlock() 260 | p.queue <- key 261 | } 262 | 263 | // markKeys marks the passed keysfor replication. 264 | func (p *peerSync) markKeys(keys []string, status replicationTask) { 265 | p.lock.Lock() 266 | for _, k := range keys { 267 | p.tasks[k] = append(p.tasks[k], status) 268 | } 269 | p.lock.Unlock() 270 | 271 | for _, k := range keys { 272 | p.queue <- k 273 | } 274 | } 275 | 276 | // replicate is a long lived process that replicates keys to the peer. 277 | func (p *peerSync) replicate() { 278 | for { 279 | select { 280 | case <-p.endSync: 281 | return 282 | case <-p.r.quit: 283 | return 284 | case key := <-p.queue: 285 | p.lock.Lock() 286 | log, ok := p.tasks[key] 287 | delete(p.tasks, key) 288 | p.lock.Unlock() 289 | if !ok { 290 | continue 291 | } 292 | 293 | task, err := collapseOpLog(log) 294 | if err != nil { 295 | p.r.n.logger.Printf("[ERROR] onecache.replicator: key %v tasks invalid: %v\n", key, err) 296 | continue 297 | } 298 | 299 | // Ensure we are still the owner of the key. 300 | if keys := p.r.n.ring.getOwnedKeys([]string{key}); len(keys) == 0 { 301 | continue 302 | } 303 | 304 | p.lock.Lock() 305 | p.replicating = key 306 | p.lock.Unlock() 307 | 308 | switch task { 309 | case task_BACKFILL: 310 | p.replicateBackfill(key, p.name) 311 | case task_DIRTY_VALUE: 312 | p.replicateDirtyValue(key, p.name) 313 | case task_TOUCH: 314 | p.replicateTouch(key, p.name) 315 | case task_DELETE: 316 | p.replicateDelete(key, p.name) 317 | } 318 | 319 | p.lock.Lock() 320 | p.replicating = "" 321 | p.lock.Unlock() 322 | } 323 | } 324 | } 325 | 326 | // replicateDirtyValue replicates a key due to a dirty value. 327 | func (p *peerSync) replicateDirtyValue(key, peer string) error { 328 | 329 | // Get the data to transmit. 330 | data, err := p.r.n.data.Get(key) 331 | if err != nil { 332 | return err 333 | } 334 | 335 | var resp struct{} 336 | sargs := StorageArgs{ 337 | Key: key, 338 | Value: data.Data, 339 | Flags: data.Flags, 340 | Exp: data.Exp, 341 | Cas: data.Cas, 342 | ExplicitPeer: peer, 343 | } 344 | v := ReplicaVersion{ 345 | Peer: p.name, 346 | Version: time.Now().UnixNano(), 347 | } 348 | args := ReplicaStorageArgs{sargs, v} 349 | 350 | if err := p.r.n.server.ReplicaSet(args, &resp); err != nil { 351 | return err 352 | } 353 | 354 | return nil 355 | } 356 | 357 | // replicateBackfill replicates a key due to leadership changes where a backfill 358 | // is necessary. 359 | func (p *peerSync) replicateBackfill(key, peer string) error { 360 | 361 | // Get the data to transmit. 362 | data, err := p.r.n.data.Get(key) 363 | if err != nil { 364 | return err 365 | } 366 | 367 | var resp struct{} 368 | args := StorageArgs{ 369 | Key: key, 370 | Value: data.Data, 371 | Flags: data.Flags, 372 | Exp: data.Exp, 373 | Cas: data.Cas, 374 | ExplicitPeer: peer, 375 | } 376 | 377 | if err := p.r.n.server.Add(args, &resp); err != nil { 378 | return err 379 | } 380 | 381 | return nil 382 | } 383 | 384 | // replicateTouch replicates a key due to it being touched. 385 | func (p *peerSync) replicateTouch(key, peer string) error { 386 | 387 | // Get the experation time. 388 | data, err := p.r.n.data.Get(key) 389 | if err != nil { 390 | return err 391 | } 392 | 393 | var resp struct{} 394 | sargs := StorageArgs{ 395 | Key: key, 396 | Exp: data.Exp, 397 | ExplicitPeer: peer, 398 | } 399 | v := ReplicaVersion{ 400 | Peer: p.name, 401 | Version: time.Now().UnixNano(), 402 | } 403 | args := ReplicaStorageArgs{sargs, v} 404 | 405 | return p.r.n.server.ReplicaTouch(args, &resp) 406 | } 407 | 408 | // replicateDelete replicates the deletion of a key. 409 | func (p *peerSync) replicateDelete(key, peer string) error { 410 | var resp struct{} 411 | args := StorageArgs{ 412 | Key: key, 413 | ExplicitPeer: peer, 414 | } 415 | 416 | return p.r.n.server.Remove(args, &resp) 417 | } 418 | 419 | // Collapses the passed operation log returning the final replication task. 420 | func collapseOpLog(log []replicationTask) (replicationTask, error) { 421 | if len(log) == 0 { 422 | return task_DIRTY_VALUE, errors.New("can not collapse empty operation log") 423 | } 424 | 425 | op := log[0] 426 | for _, s := range log[1:] { 427 | switch s { 428 | case task_DELETE: 429 | if op == task_DELETE { 430 | return task_DIRTY_VALUE, errors.New("invalid op log: two deletes in a row") 431 | } 432 | 433 | op = s 434 | case task_DIRTY_VALUE: 435 | op = s 436 | case task_TOUCH: 437 | if op == task_DELETE { 438 | return task_DIRTY_VALUE, errors.New("invalid op log: metadata can not be marked dirty after the key has been deleted") 439 | } 440 | case task_BACKFILL: 441 | op = s 442 | } 443 | } 444 | 445 | return op, nil 446 | } 447 | -------------------------------------------------------------------------------- /onecache/replicator_test.go: -------------------------------------------------------------------------------- 1 | package onecache 2 | 3 | import ( 4 | "fmt" 5 | "testing" 6 | "time" 7 | ) 8 | 9 | // mockReplicator can be used by tests to mock the replicator 10 | type mockReplicator struct { 11 | isReplicating func(key string) bool 12 | } 13 | 14 | // newMockReplicator returns a mock replicator that can be used by tests. 15 | // The passed function replaces the IsReplicating function. 16 | func newMockReplicator(isReplicating func(key string) bool) *mockReplicator { 17 | return &mockReplicator{isReplicating: isReplicating} 18 | } 19 | 20 | func (m *mockReplicator) Stop() {} 21 | func (m *mockReplicator) NodesRemoved([]string) {} 22 | func (m *mockReplicator) MarkDirty(string) {} 23 | func (m *mockReplicator) MarkTouched(key string) {} 24 | func (m *mockReplicator) MarkDeleted(key string) {} 25 | func (m *mockReplicator) MarkPeersDirty([]string) error { return nil } 26 | func (m *mockReplicator) IsReplicating(key string) bool { return m.isReplicating(key) } 27 | 28 | // Fills a node with key/value pairs with no flags and no expiration. Returns a list of the inserted data. Each entry 29 | // represents a KV pair since keys equal their values. 30 | func fillNode(t *testing.T, n *Node, numKeys int) []string { 31 | var keysAndValues []string 32 | for i := 0; i < numKeys; i++ { 33 | d := fmt.Sprintf("%d", i) 34 | keysAndValues = append(keysAndValues, d) 35 | 36 | if err := n.data.Set(d, []byte(d), 0, 0); err != nil { 37 | t.Fatalf("Set(%v) failed: %v", d, err) 38 | } 39 | } 40 | 41 | return keysAndValues 42 | } 43 | 44 | func TestMarkPeersDirtySingle(t *testing.T) { 45 | nodes, err := connectedTestNodes(2, 2) 46 | defer tearDownNodes(nodes) 47 | if err != nil { 48 | t.Fatalf("connectedTestNodes(2, 2) failed: %v", err) 49 | } 50 | 51 | n1 := nodes[0] 52 | n2 := nodes[1] 53 | 54 | // Fill node one with data. 55 | keys := fillNode(t, n1, 50) 56 | 57 | // Mark the second node as dirty. 58 | if err := n1.replicator.MarkPeersDirty([]string{n2.name}); err != nil { 59 | t.Errorf("MarkPeersDirty(%v) failed: %v", n2.name, err) 60 | } 61 | 62 | // Check all keys were replicated correctly. 63 | waitForResult(func() (bool, error) { 64 | for _, d := range n1.ring.getOwnedKeys(keys) { 65 | data, err := n2.data.Get(d) 66 | if err != nil { 67 | return false, err 68 | } 69 | 70 | act := string(data.Data) 71 | if act != d { 72 | return false, nil 73 | } 74 | } 75 | return true, nil 76 | }, func(e error) { 77 | t.Errorf("Replication failed: %v", e) 78 | }) 79 | } 80 | 81 | func TestMarkPeersDirtyRepeatedly(t *testing.T) { 82 | nodes, err := connectedTestNodes(2, 2) 83 | defer tearDownNodes(nodes) 84 | if err != nil { 85 | t.Fatalf("connectedTestNodes(2, 2) failed: %v", err) 86 | } 87 | 88 | n1 := nodes[0] 89 | n2 := nodes[1] 90 | 91 | // Fill node one with data. 92 | keys := fillNode(t, n1, 10) 93 | 94 | for i := 0; i < 5; i++ { 95 | // Mark the second node as dirty. 96 | if err := n1.replicator.MarkPeersDirty([]string{n2.name}); err != nil { 97 | t.Errorf("MarkPeersDirty(%v) failed: %v", n2.name, err) 98 | } 99 | time.Sleep(10 * time.Millisecond) 100 | } 101 | 102 | // Check all keys were replicated correctly. 103 | waitForResult(func() (bool, error) { 104 | for _, d := range n1.ring.getOwnedKeys(keys) { 105 | data, err := n2.data.Get(d) 106 | if err != nil { 107 | return false, err 108 | } 109 | 110 | act := string(data.Data) 111 | if act != d { 112 | return false, nil 113 | } 114 | } 115 | return true, nil 116 | }, func(e error) { 117 | t.Errorf("Replication failed: %v", e) 118 | }) 119 | } 120 | 121 | func TestMarkPeersDirtyMultiple(t *testing.T) { 122 | nodes, err := connectedTestNodes(4, 4) 123 | defer tearDownNodes(nodes) 124 | if err != nil { 125 | t.Fatalf("connectedTestNodes(4, 4) failed: %v", err) 126 | } 127 | 128 | n1 := nodes[0] 129 | 130 | // Fill node one with data. 131 | keys := fillNode(t, n1, 10) 132 | 133 | // Mark the other nodes as dirty. 134 | var peers []string 135 | for _, n := range nodes[1:] { 136 | peers = append(peers, n.name) 137 | } 138 | 139 | if err := n1.replicator.MarkPeersDirty(peers); err != nil { 140 | t.Errorf("MarkPeersDirty(%v) failed: %v", peers, err) 141 | } 142 | 143 | // Check all keys were replicated correctly. 144 | waitForResult(func() (bool, error) { 145 | for _, n := range nodes[1:] { 146 | for _, d := range n1.ring.getOwnedKeys(keys) { 147 | data, err := n.data.Get(d) 148 | if err != nil { 149 | return false, err 150 | } 151 | 152 | act := string(data.Data) 153 | if act != d { 154 | return false, nil 155 | } 156 | } 157 | } 158 | return true, nil 159 | }, func(e error) { 160 | t.Errorf("Replication failed: %v", e) 161 | }) 162 | } 163 | 164 | func TestMarkPeersDirtyInvalid(t *testing.T) { 165 | nodes, err := connectedTestNodes(2, 2) 166 | defer tearDownNodes(nodes) 167 | if err != nil { 168 | t.Fatalf("connectedTestNodes(2, 2) failed: %v", err) 169 | } 170 | 171 | // Fill node one with data. 172 | n1 := nodes[0] 173 | _ = fillNode(t, n1, 100) 174 | 175 | inv := "invalid" 176 | if err := nodes[0].replicator.MarkPeersDirty([]string{inv}); err == nil { 177 | t.Errorf("MarkPeersDirty(%v) should have returned an error", inv) 178 | } 179 | } 180 | 181 | // getOwner returns the owner and the original set of nodes without the owner. 182 | func getOwner(key string, nodes []*Node) (*Node, []*Node) { 183 | if len(nodes) == 0 { 184 | return nil, nil 185 | } 186 | 187 | // Get a copy of the ring and clear it. 188 | r := nodes[0].ring.ringCopy() 189 | for _, m := range r.(*consistentRing).ring.Members() { 190 | r.remove(m) 191 | } 192 | 193 | lookup := make(map[string]*Node, len(nodes)) 194 | for _, n := range nodes { 195 | lookup[n.name] = n 196 | r.add(n.name) 197 | } 198 | 199 | // Get the successor. 200 | s := r.successor(key) 201 | owner := lookup[s] 202 | 203 | var followers []*Node 204 | for name, node := range lookup { 205 | if name != s { 206 | followers = append(followers, node) 207 | } 208 | } 209 | 210 | return owner, followers 211 | } 212 | 213 | func TestMarkDirtyValue(t *testing.T) { 214 | nodes, err := connectedTestNodes(2, 2) 215 | defer tearDownNodes(nodes) 216 | if err != nil { 217 | t.Fatalf("connectedTestNodes(2, 2) failed: %v", err) 218 | } 219 | 220 | // Store a value on the owner node. 221 | k := "foo" 222 | expV := "bar" 223 | owner, followers := getOwner(k, nodes) 224 | if err := owner.data.Set(k, []byte(expV), 0, 0); err != nil { 225 | t.Fatalf("Set(%v) failed: %v", k, err) 226 | } 227 | owner.replicator.MarkDirty(k) 228 | 229 | // Wait for replication. 230 | var act string 231 | waitForResult(func() (bool, error) { 232 | data, err := followers[0].data.Get(k) 233 | if err != nil { 234 | return false, err 235 | } 236 | 237 | act = string(data.Data) 238 | if act != expV { 239 | return false, nil 240 | } 241 | return true, nil 242 | }, func(e error) { 243 | t.Errorf("Stored incorrect value; got %v; want %v", act, expV) 244 | }) 245 | } 246 | 247 | func TestMarkDirtyValueRepeated(t *testing.T) { 248 | nodes, err := connectedTestNodes(2, 2) 249 | defer tearDownNodes(nodes) 250 | if err != nil { 251 | t.Fatalf("connectedTestNodes(2, 2) failed: %v", err) 252 | } 253 | 254 | // Store a value on the owner node. 255 | k := "foo" 256 | var expV string 257 | owner, followers := getOwner(k, nodes) 258 | for i := 0; i < 20; i++ { 259 | expV = fmt.Sprintf("%d", i) 260 | if err := owner.data.Set(k, []byte(expV), 0, 0); err != nil { 261 | t.Fatalf("Set(%v) failed: %v", k, err) 262 | } 263 | owner.replicator.MarkDirty(k) 264 | } 265 | 266 | follower := followers[0] 267 | 268 | var act string 269 | waitForResult(func() (bool, error) { 270 | data, err := follower.data.Get(k) 271 | if err != nil { 272 | return false, err 273 | } 274 | 275 | act = string(data.Data) 276 | if act != expV { 277 | return false, nil 278 | } 279 | return true, nil 280 | }, func(e error) { 281 | t.Errorf("Stored incorrect value; got %v; want %v", act, expV) 282 | }) 283 | } 284 | 285 | func TestMarkDeleted(t *testing.T) { 286 | nodes, err := connectedTestNodes(2, 2) 287 | defer tearDownNodes(nodes) 288 | if err != nil { 289 | t.Fatalf("connectedTestNodes(2, 2) failed: %v", err) 290 | } 291 | 292 | // Store a value on the nodes. 293 | k := "foo" 294 | for _, node := range nodes { 295 | if err := node.data.Set(k, []byte("bar"), 0, 0); err != nil { 296 | t.Fatalf("Set(%v) failed: %v", k, err) 297 | } 298 | } 299 | 300 | owner, followers := getOwner(k, nodes) 301 | owner.replicator.MarkDeleted(k) 302 | 303 | // Check the delete happened 304 | waitForResult(func() (bool, error) { 305 | return !followers[0].data.Contains(k), nil 306 | }, func(e error) { 307 | t.Errorf("Replication failed: %v", e) 308 | }) 309 | } 310 | 311 | func TestMarkTouched(t *testing.T) { 312 | nodes, err := connectedTestNodes(2, 2) 313 | defer tearDownNodes(nodes) 314 | if err != nil { 315 | t.Fatalf("connectedTestNodes(2, 2) failed: %v", err) 316 | } 317 | 318 | // Store a key on both nodes but with different experations. 319 | k := "foo" 320 | owner, followers := getOwner(k, nodes) 321 | follower := followers[0] 322 | exp := time.Now().Add(1 * time.Hour).Unix() 323 | if err := owner.data.Set(k, []byte("bar"), exp, 0); err != nil { 324 | t.Fatalf("Set(%v) failed: %v", k, err) 325 | } 326 | 327 | if err := follower.data.Set(k, []byte("bar"), 0, 0); err != nil { 328 | t.Fatalf("Set(%v) failed: %v", k, err) 329 | } 330 | 331 | owner.replicator.MarkTouched(k) 332 | 333 | // Check the key expiration is updated. 334 | waitForResult(func() (bool, error) { 335 | data, err := follower.data.Get(k) 336 | if err != nil { 337 | return false, nil 338 | } 339 | 340 | if data.Exp != exp { 341 | return false, nil 342 | } 343 | return true, nil 344 | }, func(e error) { 345 | t.Errorf("Replication failed: %v", e) 346 | }) 347 | } 348 | 349 | func TestCollapseLogValueTouch(t *testing.T) { 350 | in := []replicationTask{task_DIRTY_VALUE, task_TOUCH} 351 | var exp replicationTask = task_DIRTY_VALUE 352 | task, err := collapseOpLog(in) 353 | if err != nil { 354 | t.Errorf("collapseOpLog(%v) returned error: %v", err) 355 | } 356 | 357 | if task != exp { 358 | t.Errorf("collapseOpLog(%v) returned %v; want", task, exp) 359 | } 360 | } 361 | 362 | func TestCollapseLogValueDelete(t *testing.T) { 363 | in := []replicationTask{task_DIRTY_VALUE, task_DELETE} 364 | var exp replicationTask = task_DELETE 365 | task, err := collapseOpLog(in) 366 | if err != nil { 367 | t.Errorf("collapseOpLog(%v) returned error: %v", err) 368 | } 369 | 370 | if task != exp { 371 | t.Errorf("collapseOpLog(%v) returned %v; want", task, exp) 372 | } 373 | } 374 | 375 | func TestCollapseLogDeleteValue(t *testing.T) { 376 | in := []replicationTask{task_DELETE, task_DIRTY_VALUE} 377 | var exp replicationTask = task_DIRTY_VALUE 378 | task, err := collapseOpLog(in) 379 | if err != nil { 380 | t.Errorf("collapseOpLog(%v) returned error: %v", err) 381 | } 382 | 383 | if task != exp { 384 | t.Errorf("collapseOpLog(%v) returned %v; want", task, exp) 385 | } 386 | } 387 | 388 | func TestCollapseLogeDeleteDelete(t *testing.T) { 389 | in := []replicationTask{task_DELETE, task_DELETE} 390 | if _, err := collapseOpLog(in); err == nil { 391 | t.Errorf("collapseOpLog(%v) should have returned an error") 392 | } 393 | } 394 | 395 | func TestCollapseLogeDeleteTouch(t *testing.T) { 396 | in := []replicationTask{task_DELETE, task_TOUCH} 397 | if _, err := collapseOpLog(in); err == nil { 398 | t.Errorf("collapseOpLog(%v) should have returned an error") 399 | } 400 | } 401 | -------------------------------------------------------------------------------- /onecache/ring.go: -------------------------------------------------------------------------------- 1 | package onecache 2 | 3 | import ( 4 | "sync" 5 | 6 | "stathat.com/c/consistent" 7 | ) 8 | 9 | // ring is an interface for a consistent hash ring that provides helpers for 10 | // accessing the owner and peers of keys. 11 | type ring interface { 12 | // add a peer to the ring. 13 | add(node string) 14 | 15 | // remove a peer from the ring. 16 | remove(node string) 17 | 18 | // replicaPeers returns the peers of a node. 19 | replicaPeers(node string) []string 20 | 21 | // successor returns the next node after the passed one in the ring. 22 | successor(node string) string 23 | 24 | // getOwnedKeys takes in a set of keys are returns the subset that is owned 25 | // by this key. 26 | getOwnedKeys(keys []string) []string 27 | 28 | // ownedKeysByNode takes in a set of keys are returns the subset that is owned 29 | // by the passed node. 30 | ownedKeysByNode(keys []string, node string) []string 31 | 32 | // ringCopy returns a copy of this node. 33 | ringCopy() ring 34 | } 35 | 36 | // consistentRing is a wrapper around a consistent hash ring that provides 37 | // helper methods for determining successors and replication peers. 38 | type consistentRing struct { 39 | self string 40 | ring *consistent.Consistent 41 | replicas int 42 | lock sync.Mutex 43 | } 44 | 45 | // newConsistentRing returns a constinent ring and takes the name of this node 46 | // and the number of its replicas. 47 | func newConsistentRing(name string, replicas int) *consistentRing { 48 | return &consistentRing{ 49 | self: name, 50 | ring: consistent.New(), 51 | replicas: replicas, 52 | } 53 | } 54 | 55 | // getOwnedKeys returns the set of keys owned by this node. 56 | func (c *consistentRing) getOwnedKeys(keys []string) []string { 57 | return c.ownedKeysByNode(keys, c.self) 58 | } 59 | 60 | // getKeysByNodes returns the set of keys owned by the passed node. 61 | func (c *consistentRing) ownedKeysByNode(keys []string, node string) []string { 62 | var owned []string 63 | for _, key := range keys { 64 | replicas := c.replicaPeers(key) 65 | if len(replicas) > 0 && replicas[0] == node { 66 | owned = append(owned, key) 67 | } 68 | } 69 | 70 | return owned 71 | } 72 | 73 | // successor returns the successor of the passed node and if there is none, an 74 | // empty string is returned. 75 | func (c *consistentRing) successor(node string) string { 76 | m := c.ring.Members() 77 | if len(m) == 1 { 78 | return m[0] 79 | } 80 | 81 | s, s2, err := c.ring.GetTwo(node) 82 | if err != nil { 83 | return "" 84 | } 85 | 86 | if s == node { 87 | return s2 88 | } 89 | 90 | return s 91 | } 92 | 93 | // replicaPeers returns the peers that this node should replicate to. 94 | func (c *consistentRing) replicaPeers(node string) []string { 95 | c.lock.Lock() 96 | defer c.lock.Unlock() 97 | 98 | peers, err := c.ring.GetN(node, c.replicas+1) 99 | if err != nil { 100 | return nil 101 | } 102 | 103 | for i, peer := range peers { 104 | if peer == node { 105 | peers = append(peers[:i], peers[i+1:]...) 106 | break 107 | } 108 | } 109 | 110 | return peers 111 | } 112 | 113 | // ringCopy returns a copy of the consistentRing. 114 | func (c *consistentRing) ringCopy() ring { 115 | c.lock.Lock() 116 | defer c.lock.Unlock() 117 | 118 | ring := &consistentRing{ 119 | self: c.self, 120 | ring: consistent.New(), 121 | replicas: c.replicas, 122 | } 123 | 124 | m := c.ring.Members() 125 | ring.ring.Set(m) 126 | return ring 127 | } 128 | 129 | // add adds a node to the consistent ring. 130 | func (c *consistentRing) add(node string) { 131 | c.lock.Lock() 132 | defer c.lock.Unlock() 133 | 134 | c.ring.Add(node) 135 | } 136 | 137 | // remove removes a node from the consistent ring. 138 | func (c *consistentRing) remove(node string) { 139 | c.lock.Lock() 140 | defer c.lock.Unlock() 141 | 142 | c.ring.Remove(node) 143 | } 144 | -------------------------------------------------------------------------------- /onecache/ring_test.go: -------------------------------------------------------------------------------- 1 | package onecache 2 | 3 | import ( 4 | "reflect" 5 | "testing" 6 | ) 7 | 8 | type mockRing struct { 9 | addImpl func(node string) 10 | removeImpl func(node string) 11 | replicaPeersImpl func(string) []string 12 | successorImpl func(string) string 13 | getOwnedKeysImpl func([]string) []string 14 | ownedKeysByNodeImpl func([]string, string) []string 15 | } 16 | 17 | func (m *mockRing) replicaPeers(node string) []string { 18 | return m.replicaPeersImpl(node) 19 | } 20 | 21 | func (m *mockRing) successor(node string) string { 22 | return m.successorImpl(node) 23 | } 24 | 25 | func (m *mockRing) getOwnedKeys(keys []string) []string { 26 | return m.getOwnedKeysImpl(keys) 27 | } 28 | 29 | func (m *mockRing) ownedKeysByNode(keys []string, node string) []string { 30 | return m.ownedKeysByNodeImpl(keys, node) 31 | } 32 | 33 | func (m *mockRing) add(node string) { m.addImpl(node) } 34 | func (m *mockRing) remove(node string) { m.removeImpl(node) } 35 | 36 | func (m *mockRing) ringCopy() ring { 37 | c := *m 38 | return &c 39 | } 40 | 41 | func TestSuccessorSingle(t *testing.T) { 42 | n, err := defaultTestNode() 43 | if err != nil { 44 | t.Fatalf("defaultTestNode() failed: %v", err) 45 | } 46 | 47 | // Get the successor and assert it is itself. 48 | s := n.ring.successor(n.name) 49 | if s != n.name { 50 | t.Errorf("successor() returned %v; want %v", s, n.name) 51 | } 52 | } 53 | 54 | func TestSuccessorValid(t *testing.T) { 55 | nodes, err := connectedTestNodes(2, 1) 56 | defer tearDownNodes(nodes) 57 | if err != nil { 58 | t.Fatalf("connectedTestNodes(2, 1) failed: %v", err) 59 | } 60 | 61 | // Get the successor and assert there is an error. 62 | n := nodes[0] 63 | s := n.ring.successor(n.name) 64 | expected := nodes[1].name 65 | if s != expected { 66 | t.Errorf("successor() returned %v; want %v", s, expected) 67 | } 68 | } 69 | 70 | func TestReplicaPeersKeySingleNode(t *testing.T) { 71 | node, err := defaultTestNode() 72 | defer node.Exit() 73 | if err != nil { 74 | t.Fatalf("connectedTestNodes(1, 0) failed: %v", err) 75 | } 76 | 77 | key := "foo" 78 | peers := node.ring.replicaPeers(key) 79 | if len(peers) != 1 { 80 | t.Errorf("replicaPeers(%v) should have one element: %v", key, peers) 81 | t.FailNow() 82 | } 83 | 84 | expected := node.name 85 | actual := peers[0] 86 | if expected != actual { 87 | t.Errorf("replicaPeers returned [%v]; expected [%v]", expected, actual) 88 | } 89 | } 90 | 91 | func TestReplicaPeersKeyMultiple(t *testing.T) { 92 | nodes, err := connectedTestNodes(2, 2) 93 | defer tearDownNodes(nodes) 94 | if err != nil { 95 | t.Fatalf("connectedTestNodes(2, 2) failed: %v", err) 96 | } 97 | 98 | n1 := nodes[0] 99 | key := "foo" 100 | peers := n1.ring.replicaPeers(key) 101 | if len(peers) != 2 { 102 | t.Errorf("replicaPeers(%v) should have two elements: %v", key, peers) 103 | t.FailNow() 104 | } 105 | } 106 | 107 | func TestReplicaPeersSelfSingle(t *testing.T) { 108 | nodes, err := connectedTestNodes(2, 2) 109 | defer tearDownNodes(nodes) 110 | if err != nil { 111 | t.Fatalf("connectedTestNodes(2, 2) failed: %v", err) 112 | } 113 | 114 | n1 := nodes[0] 115 | peers := n1.ring.replicaPeers(n1.name) 116 | if len(peers) != 1 { 117 | t.Errorf("replicaPeers(%v) should have one element: %v", n1.name, peers) 118 | t.FailNow() 119 | } 120 | 121 | expected := nodes[1].name 122 | actual := peers[0] 123 | if expected != actual { 124 | t.Errorf("replicaPeers returned [%v]; expected [%v]", expected, actual) 125 | } 126 | } 127 | 128 | func TestReplicaSelfPeers(t *testing.T) { 129 | nodes, err := connectedTestNodes(3, 2) 130 | defer tearDownNodes(nodes) 131 | if err != nil { 132 | t.Fatalf("conenctedTestNodes(3, 2) failed: %v", err) 133 | } 134 | n1 := nodes[0] 135 | 136 | peers := n1.ring.replicaPeers(n1.name) 137 | if len(peers) != 2 { 138 | t.Errorf("replicaPeers(%v) should have two elements: %v", n1.name, peers) 139 | } 140 | 141 | expected := []string{nodes[1].name, nodes[2].name} 142 | for _, e := range expected { 143 | match := false 144 | for _, a := range peers { 145 | if a == e { 146 | match = true 147 | } 148 | } 149 | 150 | if !match { 151 | t.Errorf("replicaPeers(%v) didn't include expected peer %v", n1.name, e) 152 | } 153 | } 154 | 155 | // Get successors a second time and ensure ordering is the same 156 | peers2 := n1.ring.replicaPeers(n1.name) 157 | if !reflect.DeepEqual(peers, peers2) { 158 | t.Errorf("ordering of replicaPeers(%v) should be the same over calls: %v != %v", n1.name, peers, peers2) 159 | } 160 | } 161 | -------------------------------------------------------------------------------- /onecache/rpc_dispatch.go: -------------------------------------------------------------------------------- 1 | package onecache 2 | 3 | import ( 4 | "fmt" 5 | "io" 6 | "net" 7 | "strconv" 8 | "time" 9 | ) 10 | 11 | const ( 12 | // Prepend internal RPC calls with this to distinguish from incoming memcache protocol rpc's 13 | rpcInternal byte = 0xFF 14 | 15 | // timeout for dialing another node. 16 | timeout = 10 * time.Second 17 | ) 18 | 19 | // initRpcServer intializes the listener for rpc requests on the passed port. 20 | func (n *Node) initRpcServer(port int) error { 21 | l, err := net.Listen("tcp", ":"+strconv.Itoa(port)) 22 | if err != nil { 23 | return fmt.Errorf("couldn't listen on port %d: %v", port, err) 24 | } 25 | 26 | n.port = l.Addr().(*net.TCPAddr).Port 27 | n.listener = l 28 | n.shutdown = false 29 | return nil 30 | } 31 | 32 | func (n *Node) listen() { 33 | for { 34 | conn, err := n.listener.Accept() 35 | if err != nil { 36 | select { 37 | case <-n.quitCh: 38 | return 39 | default: 40 | break 41 | } 42 | 43 | n.logger.Printf("[WARN] onecache: rpc accept failed: %v\n", err.Error()) 44 | continue 45 | } 46 | 47 | conn.(*net.TCPConn).SetNoDelay(true) 48 | conn.(*net.TCPConn).SetKeepAlive(true) 49 | go n.dispatchConn(conn) 50 | } 51 | } 52 | 53 | // dispatchConn redirects the connection to either the memcache layer or to the 54 | // internal RPC layer. 55 | func (n *Node) dispatchConn(conn net.Conn) { 56 | buf := make([]byte, 1) 57 | if _, err := conn.Read(buf); err != nil { 58 | if err != io.EOF { 59 | n.logger.Printf("[ERR] onecache: failed to read first byte: %v", err) 60 | } 61 | conn.Close() 62 | return 63 | } 64 | 65 | switch buf[0] { 66 | case rpcInternal: 67 | n.server.handleMultiplex(conn) 68 | default: 69 | n.handleMemcacheClient(conn, buf[0]) 70 | } 71 | 72 | return 73 | } 74 | 75 | // dialNode takes a peer and dails it, returning the connection for RPC use. 76 | func (n *Node) dialNode(peer string) (net.Conn, error) { 77 | n.lock.Lock() 78 | p, ok := n.nodes[peer] 79 | n.lock.Unlock() 80 | if !ok { 81 | return nil, fmt.Errorf("could not dial unknown peer %v", peer) 82 | } 83 | 84 | conn, err := net.DialTimeout("tcp", fmt.Sprintf("%s:%d", p.addr, p.port), timeout) 85 | if err != nil { 86 | return nil, err 87 | } 88 | 89 | _, err = conn.Write([]byte{rpcInternal}) 90 | if err != nil { 91 | conn.Close() 92 | return nil, err 93 | } 94 | 95 | return conn, nil 96 | } 97 | -------------------------------------------------------------------------------- /onecache/test_util.go: -------------------------------------------------------------------------------- 1 | package onecache 2 | 3 | import ( 4 | "fmt" 5 | "time" 6 | ) 7 | 8 | // Returns a test configuration suitable for unit tests. If port parameters are set to zero 9 | // an open port is chosen. 10 | func buildTestConfig(replicas, rpcPort, discoveryPort int) *NodeConfig { 11 | config := DefaultNodeConfig() 12 | config.SerfConfig.MemberlistConfig.BindAddr = "127.0.0.1" 13 | config.Replicas = replicas 14 | config.RpcPort = rpcPort 15 | config.SerfConfig.MemberlistConfig.BindPort = discoveryPort 16 | return config 17 | } 18 | 19 | // Returns a single test node with replica parameter of 2. 20 | func defaultTestNode() (*Node, error) { 21 | n, err := Create(buildTestConfig(2, 0, 0)) 22 | if err != nil { 23 | return nil, err 24 | } 25 | 26 | // Allow the node to setup. 27 | n.Start() 28 | time.Sleep(10 * time.Millisecond) 29 | 30 | return n, nil 31 | } 32 | 33 | // Returns a list of n Node's that have joined each other. 34 | func connectedTestNodes(n int, replicas int) ([]*Node, error) { 35 | configs := make([]*NodeConfig, n) 36 | for i := 0; i < n; i++ { 37 | configs[i] = buildTestConfig(replicas, 0, 0) 38 | } 39 | return buildConnectedNodes(configs) 40 | } 41 | 42 | // Returns a list of Node's created using the passed configuration and joined with each other. 43 | func buildConnectedNodes(configs []*NodeConfig) ([]*Node, error) { 44 | nodes := make([]*Node, len(configs)) 45 | for i, c := range configs { 46 | n, err := Create(c) 47 | if err != nil { 48 | return nil, err 49 | } 50 | 51 | n.Start() 52 | nodes[i] = n 53 | } 54 | 55 | // Allow the nodes to setup. 56 | time.Sleep(10 * time.Millisecond) 57 | 58 | if len(nodes) == 1 { 59 | return nodes, nil 60 | } 61 | 62 | // Create a slice containing the ports. 63 | ip := make([]string, len(nodes)) 64 | for i, n := range nodes { 65 | ip[i] = fmt.Sprintf("127.0.0.1:%d", n.serf.LocalMember().Port) 66 | } 67 | 68 | // Allow the nodes to join. 69 | for _, node := range nodes { 70 | node.Join(ip) 71 | } 72 | 73 | waitForResult(func() (bool, error) { 74 | for _, node := range nodes { 75 | if len(node.nodes) != len(nodes) { 76 | return false, nil 77 | } 78 | } 79 | return true, nil 80 | }, func(e error) {}) 81 | 82 | return nodes, nil 83 | } 84 | 85 | func tearDownNodes(nodes []*Node) { 86 | for _, n := range nodes { 87 | // Speeds up the serf leave. 88 | n.serf.Shutdown() 89 | n.Exit() 90 | } 91 | } 92 | 93 | type testFn func() (bool, error) 94 | type errorFn func(error) 95 | 96 | func waitForResult(test testFn, error errorFn) { 97 | retries := 1000 98 | 99 | for retries > 0 { 100 | time.Sleep(10 * time.Millisecond) 101 | retries-- 102 | 103 | success, err := test() 104 | if success { 105 | return 106 | } 107 | 108 | if retries == 0 { 109 | error(err) 110 | } 111 | } 112 | } 113 | -------------------------------------------------------------------------------- /ttlstore/ttl_heap.go: -------------------------------------------------------------------------------- 1 | package ttlstore 2 | 3 | import ( 4 | "container/heap" 5 | "errors" 6 | "fmt" 7 | "sync" 8 | ) 9 | 10 | type TtlHeap struct { 11 | items map[string]*Item 12 | queue minHeap 13 | lock sync.Mutex 14 | } 15 | 16 | func NewTtlHeap() *TtlHeap { 17 | return &TtlHeap{ 18 | items: make(map[string]*Item), 19 | queue: make(minHeap, 0), 20 | } 21 | } 22 | 23 | func (t *TtlHeap) Size() int { 24 | t.lock.Lock() 25 | defer t.lock.Unlock() 26 | return len(t.queue) 27 | } 28 | 29 | func (t *TtlHeap) Push(key string, time int64) error { 30 | t.lock.Lock() 31 | defer t.lock.Unlock() 32 | 33 | if _, ok := t.items[key]; ok { 34 | return fmt.Errorf("key %v already exists", key) 35 | } 36 | 37 | item := &Item{key, time, 0} 38 | t.items[key] = item 39 | t.queue.Push(item) 40 | return nil 41 | } 42 | 43 | func (t *TtlHeap) Pop() (*Item, error) { 44 | t.lock.Lock() 45 | defer t.lock.Unlock() 46 | 47 | if len(t.queue) == 0 { 48 | return nil, errors.New("heap is empty") 49 | } 50 | 51 | item := t.queue.Pop().(*Item) 52 | delete(t.items, item.Key) 53 | return item, nil 54 | } 55 | 56 | func (t *TtlHeap) Peek() (Item, error) { 57 | t.lock.Lock() 58 | defer t.lock.Unlock() 59 | 60 | if len(t.queue) == 0 { 61 | return Item{}, errors.New("heap is empty") 62 | } 63 | 64 | return *(t.queue[0]), nil 65 | } 66 | 67 | func (t *TtlHeap) Contains(key string) bool { 68 | t.lock.Lock() 69 | defer t.lock.Unlock() 70 | 71 | _, ok := t.items[key] 72 | return ok 73 | } 74 | 75 | func (t *TtlHeap) Update(key string, time int64) error { 76 | t.lock.Lock() 77 | defer t.lock.Unlock() 78 | 79 | if item, ok := t.items[key]; ok { 80 | t.queue.update(item, time) 81 | return nil 82 | } 83 | 84 | return fmt.Errorf("heap doesn't contain key %v", key) 85 | } 86 | 87 | func (t *TtlHeap) Remove(key string) error { 88 | t.lock.Lock() 89 | defer t.lock.Unlock() 90 | 91 | if item, ok := t.items[key]; ok { 92 | heap.Remove(&t.queue, item.index) 93 | delete(t.items, key) 94 | return nil 95 | } 96 | 97 | return fmt.Errorf("heap doesn't contain key %v", key) 98 | } 99 | 100 | type Item struct { 101 | Key string // The value of the Item; arbitrary. 102 | Time int64 // The priority of the Item in the queue. 103 | index int // The index of the Item in the heap. Maintained by heap operations. 104 | } 105 | 106 | // A minHeap implements heap.Interface and holds Items. 107 | type minHeap []*Item 108 | 109 | func (m minHeap) Len() int { return len(m) } 110 | 111 | func (m minHeap) Less(i, j int) bool { 112 | // We want Pop to give us the highest, not lowest, priority so we use greater than here. 113 | return m[i].Time < m[j].Time 114 | } 115 | 116 | func (m minHeap) Swap(i, j int) { 117 | m[i], m[j] = m[j], m[i] 118 | m[i].index = i 119 | m[j].index = j 120 | } 121 | 122 | func (m *minHeap) Push(x interface{}) { 123 | n := len(*m) 124 | item := x.(*Item) 125 | item.index = n 126 | *m = append(*m, item) 127 | } 128 | 129 | func (m *minHeap) Pop() interface{} { 130 | old := *m 131 | n := len(old) 132 | item := old[n-1] 133 | item.index = -1 // for safety 134 | *m = old[0 : n-1] 135 | return item 136 | } 137 | 138 | // update modifies the priority and value of an Item in the queue. 139 | func (m *minHeap) update(item *Item, time int64) { 140 | item.Time = time 141 | heap.Fix(m, item.index) 142 | } 143 | -------------------------------------------------------------------------------- /ttlstore/ttlstore.go: -------------------------------------------------------------------------------- 1 | package ttlstore 2 | 3 | import ( 4 | "errors" 5 | "fmt" 6 | "log" 7 | "math/rand" 8 | "strconv" 9 | "sync" 10 | "sync/atomic" 11 | "time" 12 | 13 | "github.com/hashicorp/go-immutable-radix" 14 | ) 15 | 16 | // DataStore is a key store that supports TTL's and LRU eviction. 17 | type DataStore struct { 18 | store *iradix.Tree 19 | rw sync.RWMutex 20 | lru *TtlHeap 21 | ttl *TtlHeap 22 | usage int64 23 | maxUsage int64 24 | rand *rand.Rand 25 | quit chan bool 26 | c *sync.Cond 27 | logger *log.Logger 28 | } 29 | 30 | type KeyData struct { 31 | Data []byte 32 | Flags int32 33 | Size int64 34 | Cas int64 35 | Exp int64 36 | } 37 | 38 | // Builder interface for KeyData 39 | func convertKeyData(i interface{}) *KeyData { 40 | return i.(*KeyData) 41 | } 42 | 43 | func newKeyData() *KeyData { 44 | return new(KeyData) 45 | } 46 | 47 | func (k *KeyData) setData(d []byte) *KeyData { 48 | k.Data = d 49 | return k 50 | } 51 | 52 | func (k *KeyData) sizeDiff(new int64) int64 { 53 | return new - k.Size 54 | } 55 | 56 | func (k *KeyData) setSize(s int64) *KeyData { 57 | k.Size = s 58 | return k 59 | } 60 | 61 | func (k *KeyData) setExp(e int64) *KeyData { 62 | k.Exp = e 63 | return k 64 | } 65 | 66 | func (k *KeyData) setFlags(f int32) *KeyData { 67 | k.Flags = f 68 | return k 69 | } 70 | 71 | func (k *KeyData) setCas(r *rand.Rand) *KeyData { 72 | k.Cas = r.Int63() 73 | return k 74 | } 75 | 76 | func New(size int64, logger *log.Logger) (*DataStore, error) { 77 | if size <= 0 { 78 | return nil, fmt.Errorf("size must be a positive value: %d", size) 79 | } 80 | 81 | var mu sync.Mutex 82 | d := &DataStore{ 83 | store: iradix.New(), 84 | lru: NewTtlHeap(), 85 | ttl: NewTtlHeap(), 86 | maxUsage: size, 87 | rand: rand.New(NewLockedSource(time.Now().UnixNano())), 88 | quit: make(chan bool), 89 | c: sync.NewCond(&mu), 90 | logger: logger, 91 | } 92 | 93 | go d.ttlWatcher() 94 | return d, nil 95 | } 96 | 97 | // Destroy removes all temp files and cleanly exits all goroutines created. 98 | func (d *DataStore) Destroy() error { 99 | close(d.quit) 100 | return nil 101 | } 102 | 103 | // ttlWatcher watches a TtlHeap for expirations. On an expiration it will reap key from either disk or memory, 104 | // determined by the passed location value. 105 | func (d *DataStore) ttlWatcher() { 106 | key := "" 107 | var exp int64 108 | var abort chan bool 109 | for { 110 | // Wait for a new key to ttl. 111 | d.c.L.Lock() 112 | for !d.change(key, exp) { 113 | select { 114 | case <-d.quit: 115 | d.c.L.Unlock() 116 | return 117 | default: 118 | } 119 | 120 | d.c.Wait() 121 | } 122 | 123 | // If it has already been reaped contine, otherwise tell the current reaper to cancel. 124 | select { 125 | case abort <- true: 126 | default: 127 | } 128 | 129 | // Update the next key to ttl. 130 | i, err := d.ttl.Peek() 131 | if err != nil { 132 | key = "" 133 | exp = 0 134 | d.c.L.Unlock() 135 | continue 136 | } 137 | key = i.Key 138 | exp = i.Time 139 | 140 | // Start the reaper. 141 | abort = make(chan bool) 142 | go d.reap(key, exp, abort) 143 | 144 | d.c.L.Unlock() 145 | } 146 | } 147 | 148 | // change returns whether the top of the ttl heap has changed from the passed 149 | // values. 150 | func (d *DataStore) change(oldKey string, oldTime int64) bool { 151 | if d.ttl.Size() == 0 { 152 | return oldKey != "" 153 | } 154 | 155 | i, _ := d.ttl.Peek() 156 | return i.Key != oldKey || i.Time != oldTime 157 | } 158 | 159 | // reap deletes the passed key after the experiation and is abortable from the 160 | // passed channel. 161 | func (d *DataStore) reap(key string, exp int64, abort <-chan bool) { 162 | expUnix := time.Unix(exp, 0) 163 | now := time.Now() 164 | dur := expUnix.Sub(now) 165 | 166 | select { 167 | case <-abort: 168 | d.logger.Printf("[DEBUG] onecache.ttlstore: aborting reaping of %v\n", key) 169 | return 170 | case <-time.After(dur): 171 | d.c.L.Lock() 172 | defer d.c.L.Unlock() 173 | defer d.c.Broadcast() 174 | if err := d.Delete(key); err != nil { 175 | d.logger.Printf("[ERROR] onecache.ttlstore: failed to reap key %v: %v\n", key, err) 176 | } 177 | } 178 | } 179 | 180 | // lruFree deletes keys until there is enough space to fit a key of the passed 181 | // size or returns an error. 182 | func (d *DataStore) lruFree(size int64) error { 183 | if d.maxUsage < size { 184 | return fmt.Errorf("memory available (%d) is less than key size %d", d.maxUsage, size) 185 | } 186 | 187 | // Evict from the LRU til there is enough space. 188 | for size > d.maxUsage-atomic.LoadInt64(&d.usage) { 189 | item, err := d.lru.Peek() 190 | if err != nil { 191 | return errors.New("can't free enough space to store in memory") 192 | } 193 | 194 | if err = d.deleteLocked(item.Key); err != nil { 195 | return err 196 | } 197 | } 198 | 199 | return nil 200 | } 201 | 202 | // List returns the set of keys stored. 203 | func (d *DataStore) List() []string { 204 | d.rw.RLock() 205 | defer d.rw.RUnlock() 206 | 207 | i := d.store.Txn().Root().Iterator() 208 | var keys []string 209 | for { 210 | key, _, ok := i.Next() 211 | if !ok { 212 | break 213 | } 214 | 215 | keys = append(keys, string(key)) 216 | } 217 | return keys 218 | } 219 | 220 | // Set stores the passed value and key metadata. 221 | func (d *DataStore) Set(key string, value []byte, exp int64, flags int32) error { 222 | d.rw.Lock() 223 | defer d.rw.Unlock() 224 | 225 | size := int64(len(value)) 226 | bKey := []byte(key) 227 | var data *KeyData 228 | var sizeUpdate int64 229 | 230 | txn := d.store.Txn() 231 | i, ok := txn.Get(bKey) 232 | if !ok { 233 | data = newKeyData().setData(value).setSize(size).setExp(exp).setFlags(flags).setCas(d.rand) 234 | sizeUpdate = size 235 | } else { 236 | data = convertKeyData(i) 237 | sizeUpdate = data.sizeDiff(size) 238 | data.setData(value).setCas(d.rand) 239 | } 240 | 241 | if err := d.lruFree(sizeUpdate); err != nil { 242 | return err 243 | } 244 | 245 | txn.Insert(bKey, data) 246 | d.store = txn.Commit() 247 | 248 | atomic.AddInt64(&d.usage, sizeUpdate) 249 | return d.touch(key, exp) 250 | } 251 | 252 | // Cas will set the key to the passed value if the CAS value is correct. 253 | func (d *DataStore) Cas(key string, value []byte, exp, cas int64, flags int32) error { 254 | d.rw.Lock() 255 | defer d.rw.Unlock() 256 | 257 | bKey := []byte(key) 258 | txn := d.store.Txn() 259 | i, ok := txn.Get(bKey) 260 | if !ok { 261 | return fmt.Errorf("cas operation invalid before key is set: %v", key) 262 | } 263 | 264 | data := convertKeyData(i) 265 | if data.Cas != cas { 266 | return errors.New("cas value did not match") 267 | } 268 | 269 | size := int64(len(value)) 270 | sizeUpdate := data.sizeDiff(size) 271 | data.setData(value).setSize(size).setExp(exp).setFlags(flags).setCas(d.rand) 272 | txn.Insert(bKey, data) 273 | 274 | if err := d.lruFree(sizeUpdate); err != nil { 275 | return err 276 | } 277 | d.store = txn.Commit() 278 | 279 | atomic.AddInt64(&d.usage, sizeUpdate) 280 | return d.touch(key, exp) 281 | } 282 | 283 | // Add stores the passed value and key metadata only if the key doesn't exist. 284 | func (d *DataStore) Add(key string, value []byte, exp int64, flags int32) error { 285 | d.rw.Lock() 286 | defer d.rw.Unlock() 287 | 288 | size := int64(len(value)) 289 | if err := d.lruFree(size); err != nil { 290 | return err 291 | } 292 | 293 | bKey := []byte(key) 294 | txn := d.store.Txn() 295 | _, ok := txn.Get(bKey) 296 | if ok { 297 | return errors.New("key is already set") 298 | } 299 | 300 | data := newKeyData().setData(value).setSize(size).setExp(exp).setFlags(flags) 301 | txn.Insert(bKey, data) 302 | d.store = txn.Commit() 303 | 304 | atomic.AddInt64(&d.usage, size) 305 | return d.touch(key, exp) 306 | } 307 | 308 | // Replace stores the passed value and key metadata only if the key exist. 309 | func (d *DataStore) Replace(key string, value []byte, exp int64, flags int32) error { 310 | d.rw.Lock() 311 | defer d.rw.Unlock() 312 | 313 | bKey := []byte(key) 314 | txn := d.store.Txn() 315 | i, ok := txn.Get(bKey) 316 | if !ok { 317 | return errors.New("key doesn't exist") 318 | } 319 | 320 | size := int64(len(value)) 321 | data := convertKeyData(i) 322 | sizeUpdate := data.sizeDiff(size) 323 | data.setData(value).setSize(size).setExp(exp).setFlags(flags) 324 | txn.Insert(bKey, data) 325 | 326 | if err := d.lruFree(sizeUpdate); err != nil { 327 | return err 328 | } 329 | d.store = txn.Commit() 330 | 331 | atomic.AddInt64(&d.usage, sizeUpdate) 332 | return d.touch(key, exp) 333 | } 334 | 335 | // Append the value to an existing value. 336 | func (d *DataStore) Append(key string, value []byte) error { 337 | d.rw.Lock() 338 | defer d.rw.Unlock() 339 | 340 | size := int64(len(value)) 341 | if err := d.lruFree(size); err != nil { 342 | return err 343 | } 344 | 345 | bKey := []byte(key) 346 | txn := d.store.Txn() 347 | i, ok := txn.Get(bKey) 348 | if !ok { 349 | return errors.New("key doesn't exist") 350 | } 351 | 352 | data := convertKeyData(i) 353 | data.setData(append(data.Data, value...)).setSize(data.Size + size) 354 | txn.Insert(bKey, data) 355 | d.store = txn.Commit() 356 | 357 | atomic.AddInt64(&d.usage, size) 358 | return d.updateLru(key) 359 | } 360 | 361 | // Prepend the value to an existing value. 362 | func (d *DataStore) Prepend(key string, value []byte) error { 363 | d.rw.Lock() 364 | defer d.rw.Unlock() 365 | 366 | size := int64(len(value)) 367 | if err := d.lruFree(size); err != nil { 368 | return err 369 | } 370 | 371 | bKey := []byte(key) 372 | txn := d.store.Txn() 373 | i, ok := txn.Get(bKey) 374 | if !ok { 375 | return errors.New("key doesn't exist") 376 | } 377 | 378 | data := convertKeyData(i) 379 | data.setData(append(value, data.Data...)).setSize(data.Size + size) 380 | txn.Insert(bKey, data) 381 | d.store = txn.Commit() 382 | 383 | atomic.AddInt64(&d.usage, size) 384 | return d.updateLru(key) 385 | } 386 | 387 | // Get returns the key value or an error if it is not set. 388 | func (d *DataStore) Get(key string) (*KeyData, error) { 389 | d.rw.RLock() 390 | defer d.rw.RUnlock() 391 | 392 | i, ok := d.store.Get([]byte(key)) 393 | if !ok { 394 | return nil, fmt.Errorf("unknown key: %v", key) 395 | } 396 | 397 | var k KeyData 398 | k = *i.(*KeyData) 399 | 400 | if err := d.updateLru(key); err != nil { 401 | return nil, err 402 | } 403 | 404 | return &k, nil 405 | } 406 | 407 | // Delete removes the key if it was previously set or returns an error. 408 | func (d *DataStore) Delete(key string) error { 409 | d.rw.Lock() 410 | defer d.rw.Unlock() 411 | return d.deleteLocked(key) 412 | } 413 | 414 | func (d *DataStore) deleteLocked(key string) error { 415 | defer d.c.Broadcast() 416 | 417 | root, i, ok := d.store.Delete([]byte(key)) 418 | d.store = root 419 | if !ok { 420 | return fmt.Errorf("key not found: %v", key) 421 | } 422 | 423 | k := i.(*KeyData) 424 | atomic.AddInt64(&d.usage, -1*k.Size) 425 | 426 | // Remove from lru and ttl. 427 | if d.ttl.Contains(key) { 428 | if err := d.ttl.Remove(key); err != nil { 429 | return fmt.Errorf("failed to remove key from ttl tracker: %v", err) 430 | } 431 | } 432 | 433 | if err := d.lru.Remove(key); err != nil { 434 | return fmt.Errorf("failed to remove key from lru tracker: %v", err) 435 | } 436 | 437 | return nil 438 | } 439 | 440 | // Increment the key if it was previously set and is a number. Returns the new 441 | // value or an error. 442 | func (d *DataStore) Increment(key string, amount int64) (int64, error) { 443 | d.rw.Lock() 444 | defer d.rw.Unlock() 445 | defer d.c.Broadcast() 446 | 447 | bKey := []byte(key) 448 | txn := d.store.Txn() 449 | i, ok := txn.Get(bKey) 450 | if !ok { 451 | return 0, errors.New("key doesn't exist") 452 | } 453 | 454 | data := convertKeyData(i) 455 | v, err := strconv.Atoi(string(data.Data)) 456 | if err != nil { 457 | return 0, fmt.Errorf("value couldn't be converted to integer: %v", v) 458 | } 459 | 460 | // Increment and convert to a string. 461 | incr := int64(v) + amount 462 | var value []byte 463 | value = strconv.AppendInt(value, incr, 10) 464 | size := int64(len(value)) 465 | sizeUpdate := data.sizeDiff(size) 466 | data.setData(value).setSize(size) 467 | 468 | if err := d.lruFree(sizeUpdate); err != nil { 469 | return 0, err 470 | } 471 | txn.Insert(bKey, data) 472 | d.store = txn.Commit() 473 | 474 | atomic.AddInt64(&d.usage, sizeUpdate) 475 | if err := d.updateLru(key); err != nil { 476 | return 0, err 477 | } 478 | 479 | return incr, nil 480 | } 481 | 482 | // Decrement the key if it was previously set and is a number, will not go below 483 | // zero. Returns the new value or an error. 484 | func (d *DataStore) Decrement(key string, amount int64) (int64, error) { 485 | d.rw.Lock() 486 | defer d.rw.Unlock() 487 | defer d.c.Broadcast() 488 | 489 | bKey := []byte(key) 490 | txn := d.store.Txn() 491 | i, ok := txn.Get(bKey) 492 | if !ok { 493 | return 0, errors.New("key doesn't exist") 494 | } 495 | 496 | data := convertKeyData(i) 497 | v, err := strconv.Atoi(string(data.Data)) 498 | if err != nil { 499 | return 0, fmt.Errorf("value couldn't be converted to integer: %v", v) 500 | } 501 | 502 | // Decrment and convert to a string. 503 | decr := int64(v) + amount 504 | if decr < 0 { 505 | decr = 0 506 | } 507 | 508 | var value []byte 509 | value = strconv.AppendInt(value, decr, 10) 510 | size := int64(len(value)) 511 | sizeUpdate := data.sizeDiff(size) 512 | data.setData(value).setSize(size) 513 | txn.Insert(bKey, data) 514 | d.store = txn.Commit() 515 | 516 | atomic.AddInt64(&d.usage, sizeUpdate) 517 | if err := d.updateLru(key); err != nil { 518 | return 0, err 519 | } 520 | 521 | return decr, nil 522 | } 523 | 524 | // Touch updates the expiration time of a key. The expiration time is in Unix 525 | // seconds, but if zero is passed the key will not expire. An error is returned 526 | // if the key does not exist. 527 | func (d *DataStore) Touch(key string, exp int64) error { 528 | d.rw.Lock() 529 | defer d.rw.Unlock() 530 | 531 | bKey := []byte(key) 532 | txn := d.store.Txn() 533 | i, ok := txn.Get(bKey) 534 | if !ok { 535 | return errors.New("key doesn't exist") 536 | } 537 | 538 | data := convertKeyData(i).setExp(exp) 539 | txn.Insert(bKey, data) 540 | d.store = txn.Commit() 541 | 542 | return d.touch(key, exp) 543 | } 544 | 545 | // touch updates the ttl and lru for the key and assumes a lock is held. 546 | func (d *DataStore) touch(key string, exp int64) error { 547 | if err := d.updateTtl(key, exp); err != nil { 548 | return err 549 | } 550 | 551 | return d.updateLru(key) 552 | } 553 | 554 | // updateTtl modifies the TtlHeap according to the expiration. Returns an error 555 | // if the key isn't found. 556 | func (d *DataStore) updateTtl(key string, exp int64) error { 557 | defer d.c.Broadcast() 558 | 559 | var err error 560 | if exp != 0 { 561 | if d.ttl.Contains(key) { 562 | err = d.ttl.Update(key, exp) 563 | } else { 564 | err = d.ttl.Push(key, exp) 565 | } 566 | } else if d.ttl.Contains(key) { 567 | err = d.ttl.Remove(key) 568 | } 569 | 570 | if err != nil { 571 | return fmt.Errorf("failed to update ttl for key %v: %v", key, err) 572 | } 573 | 574 | return nil 575 | } 576 | 577 | // updateLru moves the passed key to the top of the lru. 578 | func (d *DataStore) updateLru(key string) error { 579 | now := time.Now().Unix() 580 | var err error 581 | if d.lru.Contains(key) { 582 | err = d.lru.Update(key, now) 583 | } else { 584 | err = d.lru.Push(key, now) 585 | } 586 | 587 | if err != nil { 588 | return fmt.Errorf("failed to update lru for key %v: %v", key, err) 589 | } 590 | 591 | return nil 592 | } 593 | 594 | // Contains returns whether the key is stored. 595 | func (d *DataStore) Contains(key string) bool { 596 | d.rw.RLock() 597 | defer d.rw.RUnlock() 598 | 599 | _, ok := d.store.Get([]byte(key)) 600 | return ok 601 | } 602 | -------------------------------------------------------------------------------- /ttlstore/ttlstore_test.go: -------------------------------------------------------------------------------- 1 | package ttlstore 2 | 3 | import ( 4 | "log" 5 | "os" 6 | "reflect" 7 | "testing" 8 | "time" 9 | ) 10 | 11 | const ( 12 | testKey = "foo" 13 | testFlags = 1 14 | testNoExpire int64 = 0 15 | ) 16 | 17 | var ( 18 | testValue = []byte("bar") 19 | testSize = int64(len(testValue)) 20 | ) 21 | 22 | var ( 23 | logger = log.New(os.Stderr, "", log.LstdFlags) 24 | ) 25 | 26 | func TestNew(t *testing.T) { 27 | t.Parallel() 28 | var max int64 = 100 29 | if _, err := New(max, logger); err != nil { 30 | t.Error("New(%v) failed: %v", max, err) 31 | } 32 | 33 | max = -100 34 | if _, err := New(max, logger); err == nil { 35 | t.Error("New(%v) should have failed: %v", max, err) 36 | } 37 | } 38 | 39 | func testDataStore(t testing.TB, max int64) *DataStore { 40 | d, err := New(max, logger) 41 | if err != nil { 42 | t.Fatal("New(%v) failed: %v", max, err) 43 | } 44 | return d 45 | } 46 | 47 | func TestGetInvalidKey(t *testing.T) { 48 | t.Parallel() 49 | // Get an invalid key. 50 | d := testDataStore(t, 500) 51 | defer d.Destroy() 52 | k := "foo" 53 | if _, e := d.Get(k); e == nil { 54 | t.Errorf("Get(%v) of an invalid key should have failed", k) 55 | } 56 | } 57 | 58 | func TestSetAndGet(t *testing.T) { 59 | t.Parallel() 60 | d := testDataStore(t, 500) 61 | defer d.Destroy() 62 | if err := d.Set(testKey, testValue, testNoExpire, testFlags); err != nil { 63 | t.Errorf("Set(%v, %v, %v, %v) failed: %v", testKey, testValue, testNoExpire, testFlags, err) 64 | } 65 | 66 | data, err := d.Get(testKey) 67 | if err != nil { 68 | t.Errorf("Get(%v) failed: %v", testKey, err) 69 | } 70 | 71 | // Compare expected to actual. 72 | if !reflect.DeepEqual(data.Data, testValue) { 73 | t.Errorf("Get(%v) returned key %v; want %v", testKey, data.Data, testValue) 74 | } 75 | 76 | if data.Size != testSize { 77 | t.Errorf("Get(%v) returned size %v; want %v", testKey, data.Size, testSize) 78 | } 79 | 80 | if data.Flags != testFlags { 81 | t.Errorf("Get(%v) returned flags %v; want %v", testKey, data.Flags, testFlags) 82 | } 83 | } 84 | 85 | func TestSetLRU(t *testing.T) { 86 | t.Parallel() 87 | // Only allocate 5 bytes of memory so that two inserts force an LRU eviction. 88 | d := testDataStore(t, 5) 89 | defer d.Destroy() 90 | if err := d.Set(testKey, testValue, testNoExpire, testFlags); err != nil { 91 | t.Fatalf("Set(%v, %v, %v, %v) failed: %v", testKey, testValue, testNoExpire, testFlags, err) 92 | } 93 | 94 | // Check the key is in memory. 95 | if _, err := d.Get(testKey); err != nil { 96 | t.Errorf("Get(%v) failed: %v", testKey, err) 97 | } 98 | 99 | // Add the second key and expect the first to be LRU evicted. 100 | k2 := "bar" 101 | expV := []byte("baz") 102 | expS := int64(len(expV)) 103 | var expF int32 = 2 104 | if err := d.Set(k2, expV, testNoExpire, expF); err != nil { 105 | t.Errorf("Set(%v, %v, %v, %v) failed: %v", k2, expV, testNoExpire, expF, err) 106 | } 107 | 108 | // Check the first key is not in memory but the second is. 109 | if _, err := d.Get(testKey); err != nil { 110 | t.Error("first key still in stored") 111 | } 112 | 113 | if _, err := d.Get(k2); err != nil { 114 | t.Error("second key not stored") 115 | } 116 | 117 | data, err := d.Get(k2) 118 | if err != nil { 119 | t.Errorf("Get(%v) failed: %v", k2, err) 120 | } 121 | 122 | // Compare expected to actual. 123 | if !reflect.DeepEqual(data.Data, expV) { 124 | t.Errorf("Get(%v) returned key %v; want %v", k2, data.Data, expV) 125 | } 126 | 127 | if data.Size != expS { 128 | t.Errorf("Get(%v) returned size %v; want %v", k2, data.Size, expS) 129 | } 130 | 131 | if data.Flags != expF { 132 | t.Errorf("Get(%v) returned flags %v; want %v", k2, data.Flags, expF) 133 | } 134 | } 135 | 136 | func TestSetTooLarge(t *testing.T) { 137 | t.Parallel() 138 | d := testDataStore(t, 5) 139 | defer d.Destroy() 140 | v := []byte("this is way too large") 141 | if err := d.Set(testKey, v, testNoExpire, testFlags); err == nil { 142 | t.Errorf("Set(%v, %v, %v, %v) should have failed for being too large", testKey, v, testNoExpire, testFlags) 143 | } 144 | } 145 | 146 | func TestCasValid(t *testing.T) { 147 | t.Parallel() 148 | d := testDataStore(t, 500) 149 | defer d.Destroy() 150 | if err := d.Set(testKey, testValue, testNoExpire, testFlags); err != nil { 151 | t.Fatalf("Set(%v, %v, %v, %v) failed: %v", testKey, testValue, testNoExpire, testFlags, err) 152 | } 153 | 154 | data, err := d.Get(testKey) 155 | if err != nil { 156 | t.Fatalf("Get(%v) failed: %v", testKey, err) 157 | } 158 | 159 | expV := []byte("cas") 160 | cas := data.Cas 161 | if err := d.Cas(testKey, expV, testNoExpire, cas, testFlags); err != nil { 162 | t.Errorf("Cas(%v, %v, %v, %v, %v) failed: %v", testKey, expV, testNoExpire, cas, testFlags, err) 163 | } 164 | 165 | data, err = d.Get(testKey) 166 | if err != nil { 167 | t.Errorf("Get(%v) failed: %v", testKey, err) 168 | } 169 | 170 | // Compare epected to actual. 171 | if !reflect.DeepEqual(data.Data, expV) { 172 | t.Errorf("Get(%v) returned key %v; want %v", testKey, data.Data, expV) 173 | } 174 | 175 | if data.Cas == cas { 176 | t.Errorf("Get(%v) returned same cas value %v", testKey, cas) 177 | } 178 | } 179 | 180 | func TestCasInvalid(t *testing.T) { 181 | t.Parallel() 182 | d := testDataStore(t, 500) 183 | defer d.Destroy() 184 | if err := d.Set(testKey, testValue, testNoExpire, testFlags); err != nil { 185 | t.Fatalf("Set(%v, %v, %v, %v) failed: %v", testKey, testValue, testNoExpire, testFlags, err) 186 | } 187 | 188 | data, err := d.Get(testKey) 189 | if err != nil { 190 | t.Fatalf("Get(%v) failed: %v", testKey, err) 191 | } 192 | 193 | expV := []byte("cas") 194 | cas := data.Cas - 1 195 | if err := d.Cas(testKey, expV, testNoExpire, cas, testFlags); err == nil { 196 | t.Errorf("Cas(%v, %v, %v, %v, %v) with incorrect cas should have failed", testKey, expV, testNoExpire, cas, testFlags) 197 | } 198 | } 199 | 200 | func TestContainsMemory(t *testing.T) { 201 | t.Parallel() 202 | // Force writes to go to memory. 203 | d := testDataStore(t, 50) 204 | defer d.Destroy() 205 | if err := d.Set(testKey, testValue, testNoExpire, testFlags); err != nil { 206 | t.Fatalf("Set(%v, %v, %v, %v) failed: %v", testKey, testValue, testNoExpire, testFlags, err) 207 | } 208 | 209 | if c := d.Contains(testKey); !c { 210 | t.Errorf("Contains(%v) should return true as it is on disk", testKey) 211 | } 212 | } 213 | 214 | func TestContainsEmpty(t *testing.T) { 215 | t.Parallel() 216 | d := testDataStore(t, 500) 217 | defer d.Destroy() 218 | if d.Contains(testKey) { 219 | t.Errorf("Contains(%v) on empty DataStore returned true", testKey) 220 | } 221 | } 222 | 223 | func TestRemoveEmpty(t *testing.T) { 224 | t.Parallel() 225 | d := testDataStore(t, 500) 226 | defer d.Destroy() 227 | if err := d.Delete(testKey); err == nil { 228 | t.Errorf("Remove(%v) on empty DataStore should return error", testKey) 229 | } 230 | } 231 | 232 | func TestRemove(t *testing.T) { 233 | t.Parallel() 234 | // Force writes to go into memory. 235 | d := testDataStore(t, 50) 236 | defer d.Destroy() 237 | if err := d.Set(testKey, testValue, testNoExpire, testFlags); err != nil { 238 | t.Fatalf("Set(%v, %v, %v, %v) failed: %v", testKey, testValue, testNoExpire, testFlags, err) 239 | } 240 | 241 | if c := d.Contains(testKey); !c { 242 | t.Fatalf("Contains(%v) should return true as it is in memory", testKey) 243 | } 244 | 245 | if err := d.Delete(testKey); err != nil { 246 | t.Errorf("Remove(%v) failed: %v", testKey, err) 247 | } 248 | 249 | if c := d.Contains(testKey); c { 250 | t.Errorf("Contains(%v) should return false as the key has been removed", testKey) 251 | } 252 | } 253 | 254 | func TestTtl(t *testing.T) { 255 | t.Parallel() 256 | // Force writes to go into memory. 257 | d := testDataStore(t, 50) 258 | defer d.Destroy() 259 | e := time.Now().Add(1 * time.Second).Unix() 260 | if err := d.Set(testKey, testValue, e, testFlags); err != nil { 261 | t.Fatalf("Set(%v, %v, %v, %v) failed: %v", testKey, testValue, e, testFlags, err) 262 | } 263 | 264 | if c := d.Contains(testKey); !c { 265 | t.Fatalf("Contains(%v) should return true as it is in memory", testKey) 266 | } 267 | 268 | // Sleep pass the TTL time. 269 | time.Sleep(1500 * time.Millisecond) 270 | 271 | if c := d.Contains(testKey); c { 272 | t.Errorf("Contains(%v) should return false as the key should be TTL'd", testKey) 273 | } 274 | } 275 | 276 | func TestTouchInvalidTtl(t *testing.T) { 277 | t.Parallel() 278 | d := testDataStore(t, 50) 279 | defer d.Destroy() 280 | e := time.Now().Add(2 * time.Second).Unix() 281 | if err := d.Touch(testKey, e); err == nil { 282 | t.Errorf("Touch(%v, %v) should have failed as key does not exist", testKey, e) 283 | } 284 | } 285 | 286 | func TestTouchValidTtl(t *testing.T) { 287 | t.Parallel() 288 | // Force writes to go into memory. 289 | d := testDataStore(t, 50) 290 | defer d.Destroy() 291 | now := time.Now() 292 | e := now.Add(2 * time.Second).Unix() 293 | e2 := now.Add(4 * time.Second).Unix() 294 | if err := d.Set(testKey, testValue, e, testFlags); err != nil { 295 | t.Fatalf("Set(%v, %v, %v, %v) failed: %v", testKey, testValue, e, testFlags, err) 296 | } 297 | 298 | time.Sleep(500 * time.Millisecond) 299 | 300 | if err := d.Touch(testKey, e2); err != nil { 301 | t.Errorf("Touch(%v, %v) failed: %v", testKey, e2, err) 302 | } 303 | 304 | // Wake up at the original expiration and assert it is still there. 305 | time.Sleep(1800 * time.Millisecond) 306 | if c := d.Contains(testKey); !c { 307 | t.Errorf("Contains(%v) should be true as the key was touched", testKey) 308 | } 309 | 310 | // Wake up past the touch expiration time. 311 | time.Sleep(2 * time.Second) 312 | if c := d.Contains(testKey); c { 313 | t.Errorf("Contains(%v) should be false as the key was TTL'd", testKey) 314 | } 315 | } 316 | 317 | func TestTtlMultiple(t *testing.T) { 318 | t.Parallel() 319 | // Force writes to go into memory. 320 | d := testDataStore(t, 50) 321 | defer d.Destroy() 322 | now := time.Now() 323 | e := now.Add(1 * time.Second).Unix() 324 | e2 := now.Add(2 * time.Second).Unix() 325 | if err := d.Set(testKey, testValue, e, testFlags); err != nil { 326 | t.Fatalf("Set(%v, %v, %v, %v) failed: %v", testKey, testValue, e, testFlags, err) 327 | } 328 | 329 | keyTwo := "keyTwo" 330 | if err := d.Set(keyTwo, testValue, e2, testFlags); err != nil { 331 | t.Fatalf("Set(%v, %v, %v, %v) failed: %v", testKey, testValue, e, testFlags, err) 332 | } 333 | 334 | // Wake up after both should have been ttl'd. 335 | time.Sleep(3 * time.Second) 336 | if c := d.Contains(testKey); c { 337 | t.Errorf("Contains(%v) should be false as the key has expired", testKey) 338 | } 339 | 340 | if c := d.Contains(keyTwo); c { 341 | t.Errorf("Contains(%v) should be false as the key has expired", keyTwo) 342 | } 343 | } 344 | 345 | func TestTtlAbort(t *testing.T) { 346 | t.Parallel() 347 | // Force writes to go into memory. 348 | d := testDataStore(t, 50) 349 | defer d.Destroy() 350 | now := time.Now() 351 | e := now.Add(1 * time.Second).Unix() 352 | if err := d.Set(testKey, testValue, e, testFlags); err != nil { 353 | t.Fatalf("Set(%v, %v, %v, %v) failed: %v", testKey, testValue, e, testFlags, err) 354 | } 355 | 356 | // Touch the key to abort. 357 | if err := d.Touch(testKey, 0); err != nil { 358 | t.Fatalf("Touch(%v, 0) failed: %v", testKey, err) 359 | } 360 | 361 | // Wake up after both should have been ttl'd. 362 | time.Sleep(2 * time.Second) 363 | if c := d.Contains(testKey); !c { 364 | t.Errorf("Contains(%v) should be true as the key was touched", testKey) 365 | } 366 | } 367 | -------------------------------------------------------------------------------- /ttlstore/util.go: -------------------------------------------------------------------------------- 1 | package ttlstore 2 | 3 | import ( 4 | "math/rand" 5 | "sync" 6 | ) 7 | 8 | type lockedSource struct { 9 | lk sync.Mutex 10 | src rand.Source 11 | } 12 | 13 | func NewLockedSource(seed int64) *lockedSource { 14 | return &lockedSource{src: rand.NewSource(seed)} 15 | } 16 | 17 | func (r *lockedSource) Int63() (n int64) { 18 | r.lk.Lock() 19 | n = r.src.Int63() 20 | r.lk.Unlock() 21 | return 22 | } 23 | 24 | func (r *lockedSource) Seed(seed int64) { 25 | r.lk.Lock() 26 | r.src.Seed(seed) 27 | r.lk.Unlock() 28 | } 29 | --------------------------------------------------------------------------------