├── NOTICE ├── connstack.go ├── CONTRIBUTING.md ├── README.md ├── connstack_test.go ├── CODE_OF_CONDUCT.md ├── LICENSE ├── memcached.go └── memcached_test.go /NOTICE: -------------------------------------------------------------------------------- 1 | Comcast.github.io 2 | Copyright 2017 Comcast Cable Communications Management, LLC 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | 16 | This product includes software developed at Comcast (https://www.comcast.com/). -------------------------------------------------------------------------------- /connstack.go: -------------------------------------------------------------------------------- 1 | package memcached 2 | 3 | import ( 4 | "sync" 5 | ) 6 | 7 | // connStack is thread-safe stack of connections. 8 | type connStack struct { 9 | data []memConn 10 | mu sync.Mutex 11 | index int 12 | } 13 | 14 | // newStack instantiates new stack. 15 | func newStack(size int) *connStack { 16 | return &connStack{ 17 | data: make([]memConn, size), 18 | index: 0, 19 | } 20 | } 21 | 22 | // push pushes a connection to stack and returns false in case the stack is full. 23 | func (s *connStack) push(value memConn) bool { 24 | s.mu.Lock() 25 | defer s.mu.Unlock() 26 | if s.index >= len(s.data) { 27 | return false 28 | } 29 | s.data[s.index] = value 30 | s.index = s.index + 1 31 | return true 32 | } 33 | 34 | // pop retrieves a connection from stack and returns false in case the stack is empty. 35 | func (s *connStack) pop() (memConn, bool) { 36 | s.mu.Lock() 37 | defer s.mu.Unlock() 38 | if s.index == 0 { 39 | return nil, false 40 | } 41 | s.index = s.index - 1 42 | return s.data[s.index], true 43 | } 44 | 45 | func (s *connStack) len() int { 46 | s.mu.Lock() 47 | defer s.mu.Unlock() 48 | return s.index 49 | } 50 | -------------------------------------------------------------------------------- /CONTRIBUTING.md: -------------------------------------------------------------------------------- 1 | # How to Contribute 2 | 3 | Thank you for considering helping out! 4 | 5 | Looking to update information or feature your project? Feel free to make the change yourself. [This is the right place to get started](https://github.com/Comcast). 6 | 7 | First, please read the [code of conduct](https://github.com/Comcast/Comcast.github.io/blob/main/CODE_OF_CONDUCT.md). We take it very seriously! 8 | 9 | Next, if you want to help out, do the following: 10 | 11 | 1. Fork the project. 12 | 2. Make a feature branch with the name of your change. 13 | 3. Make a change. 14 | 4. Commit your code. 15 | 5. Issue a Pull Request. 16 | 17 | Once your PR is issued, we will review your work and decide on adding it to the project! 18 | 19 | For more details about contributing to GitHub projects see [How to use Github: Fork, Branch, Track, Squash and Pull Request](http://gun.io/blog/how-to-github-fork-branch-and-pull-request/) 20 | 21 | ## Contributor License Agreement 22 | 23 | Before Comcast merges your code into the project you must sign the [Comcast Contributor License Agreement (CLA)](https://gist.github.com/ComcastOSS/a7b8933dd8e368535378cda25c92d19a). 24 | 25 | If you haven't previously signed a Comcast CLA, you'll automatically be asked to when you open a pull request. Alternatively, we can send you a PDF that you can sign and scan back to us. Please create a new GitHub issue to request a PDF version of the CLA. -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Memcached Client Library 2 | 3 | This library implements basic operations for memcached communication such as Set and Get. It provides an idle pool of 4 | connections that can be reused as per memcached protocol the server performance is better when connections are 5 | maintained opened rather than established for every request. For applications critical to the network latency, 6 | the library provides a profiling interface to collect information about network connections, set or get times, and some 7 | detailed information for socket operations. 8 | 9 | The client instantiated with `NewClient()` is thread safe and therefore can be used in multiple goroutines. 10 | 11 | The internal pool of reusable connections is maintained for optimization purposes. The memcached protocol suggests to 12 | reuse network connections as long as possible rather than close them at the end of every request. This pool size can be 13 | set via `maxIdleConn` parameter of the `NewClient()`. In case of a peak load the new connections will still be allowed to 14 | be created even if the number exceeds the pool size. They will be closed once used so that the preferred pool size 15 | is maintained. 16 | 17 | WARNING: The default memcached server time is 2 minutes until it closes inactive connection. This library uses a default 18 | and at this point can only be changed via `Client.idleTimeout`. 19 | 20 | ## Using The Client Library 21 | 22 | This library can be used as a Go package. To create an instance of memcached client 23 | 24 | ```golang 25 | // Instantiate 26 | client, err := memcached.NewClient("127.0.0.1:11211", 100) 27 | if err != nil { 28 | return err 29 | } 30 | 31 | // Set data with key "1/2/3" and expiration time 30 seconds 32 | if err := client.Set(context.Background(), "1/2/3", []byte{"hello"}, 30); err != nil { 33 | return err 34 | } 35 | 36 | // Get data 37 | resp, err := client.Get(context.Background(), "1/2/3") 38 | if err != nil { 39 | return err 40 | } 41 | process(resp.Data) 42 | ``` 43 | 44 | To enable profiler and receive the network io values 45 | 46 | ```golang 47 | p := memcached.NewFuncProfiler(func(name string, duration time.Duration) { 48 | // put any custom code here that can read metric name/duration 49 | metrics.Record("memcached_profiler", duration, name) 50 | }) 51 | memcachedClient.SetProfiler(p) 52 | ``` 53 | 54 | To read some client stats 55 | 56 | ```golang 57 | 58 | t := time.NewTicker(1 * time.Second) 59 | for { 60 | select { 61 | case <-c.Done(): 62 | return 63 | case _ = <-t.C: 64 | stats := client.GetStats() 65 | metrics.Save("first_metric", stats.ConnPoolLen, "conn_pool_len") 66 | } 67 | } 68 | ``` 69 | 70 | ## License 71 | Comcast.github.io is licensed under [Apache License 2.0](LICENSE). Valid-License-Identifier: Apache-2.0 72 | 73 | ## Code of Conduct 74 | We take our [code of conduct](CODE_OF_CONDUCT.md) very seriously. Please abide by it. 75 | 76 | ## Contributing 77 | Please read our [contributing guide](CONTRIBUTING.md) for details on how to contribute to our project. -------------------------------------------------------------------------------- /connstack_test.go: -------------------------------------------------------------------------------- 1 | package memcached 2 | 3 | import ( 4 | "testing" 5 | "time" 6 | ) 7 | 8 | func TestPushPop(t *testing.T) { 9 | const stackSize = 10 10 | cs := newStack(stackSize) 11 | t1, _ := time.Parse(time.RFC3339, "2006-01-02T15:04:05Z") 12 | for i := 1; i <= stackSize; i++ { 13 | testConn := mockConn{} 14 | testConn.SetLastUsed(t1.Add(time.Duration(i) * time.Minute)) 15 | if ok := cs.push(&testConn); !ok { 16 | t.Fatal("stack push error: stack full") 17 | } 18 | } 19 | for i := stackSize; i > 0; i-- { 20 | testConn, ok := cs.pop() 21 | if !ok { 22 | t.Fatal("stack pop error: stack empty") 23 | } 24 | tm := t1.Add(time.Duration(i) * time.Minute) 25 | if testConn.GetLastUsed() != tm { 26 | t.Fatalf("expected time %v but got %v\n", tm, testConn.GetLastUsed()) 27 | } 28 | } 29 | } 30 | 31 | func TestPushPopError(t *testing.T) { 32 | const stackSize = 5 33 | cs := newStack(stackSize) 34 | for i := 0; i < stackSize; i++ { 35 | testConn := mockConn{} 36 | if ok := cs.push(&testConn); !ok { 37 | t.Fatal("stack push error: stack full") 38 | } 39 | } 40 | testConn := mockConn{} 41 | if ok := cs.push(&testConn); ok { 42 | t.Fatal("expected stack full") 43 | } 44 | for i := 0; i < stackSize; i++ { 45 | if _, ok := cs.pop(); !ok { 46 | t.Fatal("stack pop error: stack empty") 47 | } 48 | } 49 | if _, ok := cs.pop(); ok { 50 | t.Fatal("expected stack empty") 51 | } 52 | } 53 | 54 | func TestPopAfter(t *testing.T) { 55 | const stackSize = 10 56 | 57 | t1, _ := time.Parse(time.RFC3339, "2006-01-02T15:04:05Z") 58 | 59 | type test struct { 60 | name string 61 | time time.Time 62 | expectedTimeAfter bool 63 | } 64 | 65 | tests := []test{ 66 | { 67 | name: "Connection not expired", 68 | time: t1.Add(1 * time.Minute), 69 | expectedTimeAfter: true, 70 | }, 71 | { 72 | name: "Connection not expired", 73 | time: t1.Add(2 * time.Minute), 74 | expectedTimeAfter: true, 75 | }, 76 | { 77 | name: "Connection not expired", 78 | time: t1.Add(3 * time.Minute), 79 | expectedTimeAfter: true, 80 | }, 81 | { 82 | name: "Connection expired", 83 | time: t1.Add(time.Duration(stackSize+1) * time.Minute), 84 | expectedTimeAfter: false, 85 | }, 86 | { 87 | name: "Connection not expired", 88 | time: t1.Add(4 * time.Minute), 89 | expectedTimeAfter: true, 90 | }, 91 | } 92 | 93 | for _, test := range tests { 94 | t.Run(test.name, func(t *testing.T) { 95 | cs := newStack(stackSize) 96 | for i := 0; i < stackSize; i++ { 97 | testConn := mockConn{} 98 | testConn.SetLastUsed(t1.Add(time.Duration(i) * time.Minute)) 99 | if ok := cs.push(&testConn); !ok { 100 | t.Fatal("stack push error: stack full") 101 | } 102 | } 103 | conn, ok := cs.pop() 104 | if !ok { 105 | t.Fatal("connection stack is empty") 106 | } 107 | ta := conn.GetLastUsed().After(test.time) 108 | if test.expectedTimeAfter != ta { 109 | t.Fatalf("expected time after %v but got %v", test.expectedTimeAfter, ta) 110 | } 111 | }) 112 | } 113 | } 114 | -------------------------------------------------------------------------------- /CODE_OF_CONDUCT.md: -------------------------------------------------------------------------------- 1 | # Contributor Covenant Code of Conduct 2 | 3 | ## Our Pledge 4 | 5 | In the interest of fostering an open and welcoming environment, we as 6 | contributors and maintainers pledge to making participation in our project and 7 | our community a harassment-free experience for everyone, regardless of age, body 8 | size, disability, ethnicity, gender identity and expression, level of experience, 9 | nationality, personal appearance, race, religion, or sexual identity and 10 | orientation. 11 | 12 | ## Our Standards 13 | 14 | Examples of behavior that contributes to creating a positive environment 15 | include: 16 | 17 | * Using welcoming and inclusive language 18 | * Being respectful of differing viewpoints and experiences 19 | * Gracefully accepting constructive criticism 20 | * Focusing on what is best for the community 21 | * Showing empathy towards other community members 22 | 23 | Examples of unacceptable behavior by participants include: 24 | 25 | * The use of sexualized language or imagery and unwelcome sexual attention or 26 | advances 27 | * Trolling, insulting/derogatory comments, and personal or political attacks 28 | * Public or private harassment 29 | * Publishing others' private information, such as a physical or electronic 30 | address, without explicit permission 31 | * Other conduct which could reasonably be considered inappropriate in a 32 | professional setting 33 | 34 | ## Our Responsibilities 35 | 36 | Project maintainers are responsible for clarifying the standards of acceptable 37 | behavior and are expected to take appropriate and fair corrective action in 38 | response to any instances of unacceptable behavior. 39 | 40 | Project maintainers have the right and responsibility to remove, edit, or 41 | reject comments, commits, code, wiki edits, issues, and other contributions 42 | that are not aligned to this Code of Conduct, or to ban temporarily or 43 | permanently any contributor for other behaviors that they deem inappropriate, 44 | threatening, offensive, or harmful. 45 | 46 | ## Scope 47 | 48 | This Code of Conduct applies both within project spaces and in public spaces 49 | when an individual is representing the project or its community. Examples of 50 | representing a project or community include using an official project e-mail 51 | address, posting via an official social media account, or acting as an appointed 52 | representative at an online or offline event. Representation of a project may be 53 | further defined and clarified by project maintainers. 54 | 55 | ## Enforcement 56 | 57 | Instances of abusive, harassing, or otherwise unacceptable behavior may be 58 | reported by contacting the project team at [Comcast_Open_Source_Services@comcast.com](mailto:Comcast_Open_Source_Services@comcast.com). 59 | All complaints will be reviewed and investigated and will result in a response that 60 | is deemed necessary and appropriate to the circumstances. The project team is 61 | obligated to maintain confidentiality with regard to the reporter of an incident. 62 | Further details of specific enforcement policies may be posted separately. 63 | 64 | Project maintainers who do not follow or enforce the Code of Conduct in good 65 | faith may face temporary or permanent repercussions as determined by other 66 | members of the project's leadership. 67 | 68 | ## Attribution 69 | 70 | This Code of Conduct is adapted from the [Contributor Covenant][homepage], version 1.4, 71 | available at [contributor-covenant.org/version/1/4][version] 72 | 73 | [homepage]: http://contributor-covenant.org 74 | [version]: http://contributor-covenant.org/version/1/4/ -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Apache License 2 | Version 2.0, January 2004 3 | http://www.apache.org/licenses/ 4 | 5 | TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 6 | 7 | 1. Definitions. 8 | 9 | "License" shall mean the terms and conditions for use, reproduction, 10 | and distribution as defined by Sections 1 through 9 of this document. 11 | 12 | "Licensor" shall mean the copyright owner or entity authorized by 13 | the copyright owner that is granting the License. 14 | 15 | "Legal Entity" shall mean the union of the acting entity and all 16 | other entities that control, are controlled by, or are under common 17 | control with that entity. For the purposes of this definition, 18 | "control" means (i) the power, direct or indirect, to cause the 19 | direction or management of such entity, whether by contract or 20 | otherwise, or (ii) ownership of fifty percent (50%) or more of the 21 | outstanding shares, or (iii) beneficial ownership of such entity. 22 | 23 | "You" (or "Your") shall mean an individual or Legal Entity 24 | exercising permissions granted by this License. 25 | 26 | "Source" form shall mean the preferred form for making modifications, 27 | including but not limited to software source code, documentation 28 | source, and configuration files. 29 | 30 | "Object" form shall mean any form resulting from mechanical 31 | transformation or translation of a Source form, including but 32 | not limited to compiled object code, generated documentation, 33 | and conversions to other media types. 34 | 35 | "Work" shall mean the work of authorship, whether in Source or 36 | Object form, made available under the License, as indicated by a 37 | copyright notice that is included in or attached to the work 38 | (an example is provided in the Appendix below). 39 | 40 | "Derivative Works" shall mean any work, whether in Source or Object 41 | form, that is based on (or derived from) the Work and for which the 42 | editorial revisions, annotations, elaborations, or other modifications 43 | represent, as a whole, an original work of authorship. For the purposes 44 | of this License, Derivative Works shall not include works that remain 45 | separable from, or merely link (or bind by name) to the interfaces of, 46 | the Work and Derivative Works thereof. 47 | 48 | "Contribution" shall mean any work of authorship, including 49 | the original version of the Work and any modifications or additions 50 | to that Work or Derivative Works thereof, that is intentionally 51 | submitted to Licensor for inclusion in the Work by the copyright owner 52 | or by an individual or Legal Entity authorized to submit on behalf of 53 | the copyright owner. For the purposes of this definition, "submitted" 54 | means any form of electronic, verbal, or written communication sent 55 | to the Licensor or its representatives, including but not limited to 56 | communication on electronic mailing lists, source code control systems, 57 | and issue tracking systems that are managed by, or on behalf of, the 58 | Licensor for the purpose of discussing and improving the Work, but 59 | excluding communication that is conspicuously marked or otherwise 60 | designated in writing by the copyright owner as "Not a Contribution." 61 | 62 | "Contributor" shall mean Licensor and any individual or Legal Entity 63 | on behalf of whom a Contribution has been received by Licensor and 64 | subsequently incorporated within the Work. 65 | 66 | 2. Grant of Copyright License. Subject to the terms and conditions of 67 | this License, each Contributor hereby grants to You a perpetual, 68 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 69 | copyright license to reproduce, prepare Derivative Works of, 70 | publicly display, publicly perform, sublicense, and distribute the 71 | Work and such Derivative Works in Source or Object form. 72 | 73 | 3. Grant of Patent License. Subject to the terms and conditions of 74 | this License, each Contributor hereby grants to You a perpetual, 75 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 76 | (except as stated in this section) patent license to make, have made, 77 | use, offer to sell, sell, import, and otherwise transfer the Work, 78 | where such license applies only to those patent claims licensable 79 | by such Contributor that are necessarily infringed by their 80 | Contribution(s) alone or by combination of their Contribution(s) 81 | with the Work to which such Contribution(s) was submitted. If You 82 | institute patent litigation against any entity (including a 83 | cross-claim or counterclaim in a lawsuit) alleging that the Work 84 | or a Contribution incorporated within the Work constitutes direct 85 | or contributory patent infringement, then any patent licenses 86 | granted to You under this License for that Work shall terminate 87 | as of the date such litigation is filed. 88 | 89 | 4. Redistribution. You may reproduce and distribute copies of the 90 | Work or Derivative Works thereof in any medium, with or without 91 | modifications, and in Source or Object form, provided that You 92 | meet the following conditions: 93 | 94 | (a) You must give any other recipients of the Work or 95 | Derivative Works a copy of this License; and 96 | 97 | (b) You must cause any modified files to carry prominent notices 98 | stating that You changed the files; and 99 | 100 | (c) You must retain, in the Source form of any Derivative Works 101 | that You distribute, all copyright, patent, trademark, and 102 | attribution notices from the Source form of the Work, 103 | excluding those notices that do not pertain to any part of 104 | the Derivative Works; and 105 | 106 | (d) If the Work includes a "NOTICE" text file as part of its 107 | distribution, then any Derivative Works that You distribute must 108 | include a readable copy of the attribution notices contained 109 | within such NOTICE file, excluding those notices that do not 110 | pertain to any part of the Derivative Works, in at least one 111 | of the following places: within a NOTICE text file distributed 112 | as part of the Derivative Works; within the Source form or 113 | documentation, if provided along with the Derivative Works; or, 114 | within a display generated by the Derivative Works, if and 115 | wherever such third-party notices normally appear. The contents 116 | of the NOTICE file are for informational purposes only and 117 | do not modify the License. You may add Your own attribution 118 | notices within Derivative Works that You distribute, alongside 119 | or as an addendum to the NOTICE text from the Work, provided 120 | that such additional attribution notices cannot be construed 121 | as modifying the License. 122 | 123 | You may add Your own copyright statement to Your modifications and 124 | may provide additional or different license terms and conditions 125 | for use, reproduction, or distribution of Your modifications, or 126 | for any such Derivative Works as a whole, provided Your use, 127 | reproduction, and distribution of the Work otherwise complies with 128 | the conditions stated in this License. 129 | 130 | 5. Submission of Contributions. Unless You explicitly state otherwise, 131 | any Contribution intentionally submitted for inclusion in the Work 132 | by You to the Licensor shall be under the terms and conditions of 133 | this License, without any additional terms or conditions. 134 | Notwithstanding the above, nothing herein shall supersede or modify 135 | the terms of any separate license agreement you may have executed 136 | with Licensor regarding such Contributions. 137 | 138 | 6. Trademarks. This License does not grant permission to use the trade 139 | names, trademarks, service marks, or product names of the Licensor, 140 | except as required for reasonable and customary use in describing the 141 | origin of the Work and reproducing the content of the NOTICE file. 142 | 143 | 7. Disclaimer of Warranty. Unless required by applicable law or 144 | agreed to in writing, Licensor provides the Work (and each 145 | Contributor provides its Contributions) on an "AS IS" BASIS, 146 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or 147 | implied, including, without limitation, any warranties or conditions 148 | of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A 149 | PARTICULAR PURPOSE. You are solely responsible for determining the 150 | appropriateness of using or redistributing the Work and assume any 151 | risks associated with Your exercise of permissions under this License. 152 | 153 | 8. Limitation of Liability. In no event and under no legal theory, 154 | whether in tort (including negligence), contract, or otherwise, 155 | unless required by applicable law (such as deliberate and grossly 156 | negligent acts) or agreed to in writing, shall any Contributor be 157 | liable to You for damages, including any direct, indirect, special, 158 | incidental, or consequential damages of any character arising as a 159 | result of this License or out of the use or inability to use the 160 | Work (including but not limited to damages for loss of goodwill, 161 | work stoppage, computer failure or malfunction, or any and all 162 | other commercial damages or losses), even if such Contributor 163 | has been advised of the possibility of such damages. 164 | 165 | 9. Accepting Warranty or Additional Liability. While redistributing 166 | the Work or Derivative Works thereof, You may choose to offer, 167 | and charge a fee for, acceptance of support, warranty, indemnity, 168 | or other liability obligations and/or rights consistent with this 169 | License. However, in accepting such obligations, You may act only 170 | on Your own behalf and on Your sole responsibility, not on behalf 171 | of any other Contributor, and only if You agree to indemnify, 172 | defend, and hold each Contributor harmless for any liability 173 | incurred by, or claims asserted against, such Contributor by reason 174 | of your accepting any such warranty or additional liability. 175 | 176 | END OF TERMS AND CONDITIONS 177 | 178 | APPENDIX: How to apply the Apache License to your work. 179 | 180 | To apply the Apache License to your work, attach the following 181 | boilerplate notice, with the fields enclosed by brackets "[]" 182 | replaced with your own identifying information. (Don't include 183 | the brackets!) The text should be enclosed in the appropriate 184 | comment syntax for the file format. We also recommend that a 185 | file or class name and description of purpose be included on the 186 | same "printed page" as the copyright notice for easier 187 | identification within third-party archives. 188 | 189 | Copyright [yyyy] [name of copyright owner] 190 | 191 | Licensed under the Apache License, Version 2.0 (the "License"); 192 | you may not use this file except in compliance with the License. 193 | You may obtain a copy of the License at 194 | 195 | http://www.apache.org/licenses/LICENSE-2.0 196 | 197 | Unless required by applicable law or agreed to in writing, software 198 | distributed under the License is distributed on an "AS IS" BASIS, 199 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 200 | See the License for the specific language governing permissions and 201 | limitations under the License. 202 | -------------------------------------------------------------------------------- /memcached.go: -------------------------------------------------------------------------------- 1 | package memcached 2 | 3 | import ( 4 | "bytes" 5 | "context" 6 | "errors" 7 | "fmt" 8 | "net" 9 | "strconv" 10 | "strings" 11 | "sync/atomic" 12 | "time" 13 | ) 14 | 15 | // defaultIdleConn is default number of idle connections. 16 | const defaultIdleConn = 20 17 | 18 | // internalBuffSize is size of a buffer used for reading network data. 19 | const internalBuffSize = 1024 20 | 21 | // defaultIdleConnectionTimeout is default value for memcached server connection timeout. 22 | const defaultIdleConnectionTimeout = 2 * time.Minute 23 | 24 | // ErrKeyNotFound returned when given key not found. 25 | var ErrKeyNotFound = errors.New("key not found") 26 | 27 | // Client is used for basic operations such as set and get. 28 | type Client struct { 29 | addr string 30 | pool *connStack 31 | provider connProvider 32 | idleTimeout time.Duration 33 | timeNow func() time.Time 34 | profiler Profiler 35 | connCounter uint32 36 | } 37 | 38 | // Response structure for operations such as set and get. 39 | type Response struct { 40 | // Key is a unique key that has some value 41 | Key string 42 | // Data is the value associated with the key 43 | Data []byte 44 | } 45 | 46 | type connProvider interface { 47 | NewConn(address string) (memConn, error) 48 | } 49 | 50 | var ( 51 | endResp = []byte{'E', 'N', 'D', '\r', '\n'} 52 | sep = []byte{'\r', '\n'} 53 | ) 54 | 55 | // NewClient instantiates a client with a given address in a form of host:port 56 | // and maxIdleConn conn pool. 57 | // Max idle connection specifies the max number of reused connections and should 58 | // be greater than zero. 59 | func NewClient(addr string, maxIdleConn int) (*Client, error) { 60 | if maxIdleConn < 1 { 61 | return nil, fmt.Errorf("maxIdleConn must be greater than zero") 62 | } 63 | return &Client{ 64 | addr: addr, 65 | pool: newStack(maxIdleConn), 66 | provider: &tcpConnProvider{}, 67 | idleTimeout: defaultIdleConnectionTimeout, 68 | timeNow: time.Now, 69 | profiler: &nopProfiler{}, 70 | }, nil 71 | } 72 | 73 | // Set sets value associated with a key. If the context contains timeout, 74 | // a deadline for the connection will be used. It is up to the client to set the timeout. 75 | // Returns an error in case NOT_STORED response received from the server or if 76 | // any other error occurred such as connection error, etc. 77 | func (c *Client) Set(ctx context.Context, key string, data []byte, expiration int) error { 78 | start := time.Now() 79 | defer func() { c.profiler.Report("set_duration", time.Since(start)) }() 80 | // preferably reuse existing conn per memcached protocol 81 | conn, err := c.acquireConn(c.addr) 82 | if err != nil { 83 | return err 84 | } 85 | defer c.releaseConn(conn) 86 | 87 | if deadline, ok := ctx.Deadline(); ok { 88 | _ = c.setConnDeadline(conn, deadline) 89 | } 90 | 91 | // data can be empty in which case length is 0 92 | var length int 93 | if data != nil { 94 | length = len(data) 95 | } 96 | 97 | // write command and data 98 | command := fmt.Sprintf("%v %v %v %v %v\r\n", "set", key, uint16(0), expiration, length) 99 | t := time.Now() 100 | if _, err := conn.Write([]byte(command)); err != nil { 101 | return fmt.Errorf("error write conn data: %v", err) 102 | } 103 | if _, err := conn.Write(data); err != nil { 104 | return fmt.Errorf("error write conn data: %v", err) 105 | } 106 | if _, err := conn.Write([]byte("\r\n")); err != nil { 107 | return fmt.Errorf("error write conn data: %v", err) 108 | } 109 | c.profiler.Report("set_write_data", time.Since(t)) 110 | 111 | // read response 112 | var mv maxValueCollector 113 | buf := make([]byte, internalBuffSize) 114 | end := 0 115 | t = time.Now() 116 | for { 117 | select { 118 | case <-ctx.Done(): 119 | return ctx.Err() 120 | default: 121 | } 122 | if internalBuffSize == end { 123 | return fmt.Errorf("error read response: internal buffer overflow") 124 | } 125 | rt := time.Now() 126 | br, err := conn.Read(buf[end:]) 127 | mv.Add(time.Since(rt)) 128 | if err != nil { 129 | return fmt.Errorf("error read response: %v", err) 130 | } 131 | end += br 132 | if bytes.Contains(buf, sep) { 133 | break 134 | } 135 | } 136 | c.profiler.Report("set_read_response", time.Since(t)) 137 | c.profiler.Report("set_read_one_max", mv.Max()) 138 | 139 | result := string(buf[:end]) 140 | 141 | switch result { 142 | case "STORED\r\n": 143 | return nil 144 | case "NOT_STORED\r\n": 145 | return fmt.Errorf("server NOT_STORED received: condition for set command isn't met") 146 | } 147 | 148 | return fmt.Errorf("set error: %v", result) 149 | } 150 | 151 | // Get returns the value associated with the key. If the context contains timeout, 152 | // a deadline for the connection will be used. It is up to the client to set the timeout. 153 | // Returns ErrKeyNotFound for non-existing key. 154 | func (c *Client) Get(ctx context.Context, key string) (*Response, error) { 155 | start := time.Now() 156 | defer func() { c.profiler.Report("get_duration", time.Since(start)) }() 157 | const ( 158 | ParseValueSection = 0 159 | ParseDataSection = 1 160 | Complete = 2 161 | ) 162 | // preferably reuse existing conn per memcached protocol 163 | conn, err := c.acquireConn(c.addr) 164 | if err != nil { 165 | return nil, fmt.Errorf("error acquire conn: %v", err) 166 | } 167 | defer c.releaseConn(conn) 168 | 169 | if deadline, ok := ctx.Deadline(); ok { 170 | _ = c.setConnDeadline(conn, deadline) 171 | } 172 | 173 | // write command and data 174 | command := fmt.Sprintf("%v %v\r\n", "get", key) 175 | t := time.Now() 176 | if _, err := conn.Write([]byte(command)); err != nil { 177 | return nil, fmt.Errorf("error write conn data: %v", err) 178 | } 179 | c.profiler.Report("get_write_data", time.Since(t)) 180 | 181 | response := Response{ 182 | Key: key, 183 | } 184 | 185 | // read response VALUE part 186 | buf := make([]byte, internalBuffSize) 187 | end := 0 // position right after the data in the result array where new data can be written 188 | dataLen := 0 // contains data section length (without \r\n and END\r\n) 189 | t = time.Now() 190 | for step := ParseValueSection; step != Complete; { 191 | select { 192 | case <-ctx.Done(): 193 | return nil, fmt.Errorf("context error: %v", ctx.Err()) 194 | default: 195 | } 196 | if internalBuffSize == end && step == ParseValueSection { 197 | return nil, fmt.Errorf("error read response: internal buffer overflow") 198 | } 199 | 200 | br, err := conn.Read(buf[end:]) 201 | if err != nil { 202 | return nil, fmt.Errorf("error read response: %v", err) 203 | } 204 | end += br 205 | 206 | if step == ParseValueSection { 207 | // analyze data - it can contain VALUE section and optionally part of the data section 208 | if bytes.Equal(buf[:end], endResp) { 209 | return nil, ErrKeyNotFound 210 | } 211 | if parts := split(buf[:end]); parts != nil { 212 | _, dataLen, err = parseValueResp(string(parts[0])) 213 | if err != nil { 214 | return nil, fmt.Errorf("error parse resp: %v", err) 215 | } 216 | buf = make([]byte, dataLen+len(sep)+len(endResp)) 217 | end = copy(buf, parts[1]) 218 | step = ParseDataSection 219 | } 220 | } 221 | if step == ParseDataSection { 222 | data := buf[:end] 223 | if len(data) < dataLen { 224 | continue 225 | } 226 | idx := bytes.LastIndex(data, endResp) 227 | if idx == -1 || idx < dataLen { 228 | continue 229 | } 230 | data = data[:idx] 231 | if bytes.HasSuffix(data, sep) { 232 | // make sure to only trim the last 'sep' 233 | response.Data = data[:len(data)-len(sep)] 234 | } else { 235 | response.Data = data 236 | } 237 | if response.Data == nil { 238 | response.Data = make([]byte, 0) 239 | } 240 | step = Complete 241 | } 242 | } 243 | c.profiler.Report("get_read_response", time.Since(t)) 244 | 245 | return &response, nil 246 | } 247 | 248 | // SetProfiler sets a profiler to measure network io read/write times. 249 | func (c *Client) SetProfiler(prof Profiler) { 250 | c.profiler = prof 251 | } 252 | 253 | // GetProfiler returns profiler. 254 | func (c *Client) GetProfiler() Profiler { 255 | return c.profiler 256 | } 257 | 258 | // GetStats returns internal stats such as number of connections, etc. 259 | func (c *Client) GetStats() Stats { 260 | return Stats{ 261 | ConnPoolLen: c.pool.len(), 262 | OpenConnCount: atomic.LoadUint32(&c.connCounter), 263 | } 264 | } 265 | 266 | func (c *Client) acquireConn(address string) (memConn, error) { 267 | t := time.Now() 268 | defer func() { c.profiler.Report("conn_acquire", time.Since(t)) }() 269 | // maxTime is time in the past which starts counting since 270 | // when a connection is still active on the server side. 271 | // If this connection was last used by the client any time 272 | // before it, the server has likely closed it by timeNow, and we 273 | // need to discard it 274 | maxTime := c.timeNow().Add(-c.idleTimeout) 275 | for { 276 | conn, ok := c.pool.pop() 277 | if !ok { 278 | return c.newConn(address) 279 | } 280 | if conn.GetLastUsed().After(maxTime) { 281 | // since we want to move to the next connection if 282 | // SetDeadline fails, unlike typical GoLang "if" 283 | // this doesn't check for err != nil 284 | if conn.SetDeadline(time.Time{}) == nil { 285 | return conn, nil 286 | } 287 | } 288 | _ = c.closeConn(conn) // ignore returned error since it is not essential to this function 289 | } 290 | } 291 | 292 | func (c *Client) releaseConn(conn memConn) error { 293 | t := time.Now() 294 | defer func() { c.profiler.Report("conn_release", time.Since(t)) }() 295 | // this connection may not be healthy. We don't want 296 | // to put it back in the pool 297 | if conn.LastErr() != nil { 298 | return c.closeConn(conn) 299 | } 300 | // set when this connection was last used 301 | // to properly close idle connections due to timeout 302 | conn.SetLastUsed(c.timeNow()) 303 | // put connection back in the pool 304 | if ok := c.pool.push(conn); !ok { 305 | return c.closeConn(conn) 306 | } 307 | return nil 308 | } 309 | 310 | func (c *Client) newConn(address string) (memConn, error) { 311 | t := time.Now() 312 | newConn, err := c.provider.NewConn(address) 313 | d := time.Since(t) 314 | if err != nil { 315 | c.profiler.Report("conn_create_new_err", d) 316 | } else { 317 | atomic.AddUint32(&c.connCounter, 1) 318 | c.profiler.Report("conn_create_new", d) 319 | } 320 | return newConn, err 321 | } 322 | 323 | func (c *Client) closeConn(conn memConn) error { 324 | atomic.AddUint32(&c.connCounter, ^uint32(0)) // idiomatic decrement - https://pkg.go.dev/sync/atomic#AddUint32 325 | t := time.Now() 326 | err := conn.Close() 327 | d := time.Since(t) 328 | if err != nil { 329 | c.profiler.Report("conn_close_err", d) 330 | } else { 331 | c.profiler.Report("conn_close", d) 332 | } 333 | return err 334 | } 335 | 336 | func (c *Client) setConnDeadline(conn memConn, deadline time.Time) error { 337 | t := time.Now() 338 | err := conn.SetDeadline(deadline) 339 | c.profiler.Report("conn_set_deadline", time.Since(t)) 340 | return err 341 | } 342 | 343 | // Stats represents internal stats such as connections count, etc. 344 | type Stats struct { 345 | ConnPoolLen int // ConnPoolLen is the current size of connections pool 346 | OpenConnCount uint32 // OpenConnCount is the current number of opened connections 347 | } 348 | 349 | type tcpConnProvider struct{} 350 | 351 | func (p *tcpConnProvider) NewConn(address string) (memConn, error) { 352 | conn, err := net.Dial("tcp", address) 353 | return &tcpConn{tcp: conn, lastErr: err}, err 354 | } 355 | 356 | // Profiler is an interface for profiling the client network io performance. 357 | type Profiler interface { 358 | Report(name string, duration time.Duration) 359 | } 360 | 361 | type nopProfiler struct{} 362 | 363 | func (n *nopProfiler) Report(name string, duration time.Duration) { 364 | } 365 | 366 | type FuncProfiler struct { 367 | callbackFunc func(string, time.Duration) 368 | } 369 | 370 | // NewFuncProfiler instantiates profiler. func f is exposed to the client of 371 | // this method to customize reported metrics handling. 372 | func NewFuncProfiler(f func(string, time.Duration)) *FuncProfiler { 373 | return &FuncProfiler{ 374 | callbackFunc: f, 375 | } 376 | } 377 | 378 | // Report reports given metric name and duration to the caller. 379 | func (f *FuncProfiler) Report(name string, duration time.Duration) { 380 | if f.callbackFunc != nil { 381 | f.callbackFunc(name, duration) 382 | } 383 | } 384 | 385 | // memConn encapsulates net.Conn interface and saves last error since net.Conn 386 | // doesn't provide any information on whether the connection is closed 387 | type memConn interface { 388 | net.Conn 389 | LastErr() error 390 | GetLastUsed() time.Time 391 | SetLastUsed(t time.Time) 392 | } 393 | 394 | type tcpConn struct { 395 | tcp net.Conn 396 | lastErr error 397 | lastUsed time.Time 398 | } 399 | 400 | func (c *tcpConn) Read(b []byte) (int, error) { 401 | var n int 402 | n, c.lastErr = c.tcp.Read(b) 403 | return n, c.lastErr 404 | } 405 | 406 | func (c *tcpConn) Write(b []byte) (int, error) { 407 | var n int 408 | n, c.lastErr = c.tcp.Write(b) 409 | return n, c.lastErr 410 | } 411 | 412 | func (c *tcpConn) Close() error { 413 | c.lastErr = c.tcp.Close() 414 | return c.lastErr 415 | } 416 | 417 | func (c *tcpConn) LocalAddr() net.Addr { 418 | return c.tcp.LocalAddr() 419 | } 420 | 421 | func (c *tcpConn) RemoteAddr() net.Addr { 422 | return c.tcp.RemoteAddr() 423 | } 424 | 425 | func (c *tcpConn) SetDeadline(t time.Time) error { 426 | c.lastErr = c.tcp.SetDeadline(t) 427 | return c.lastErr 428 | } 429 | 430 | func (c *tcpConn) SetReadDeadline(t time.Time) error { 431 | c.lastErr = c.tcp.SetReadDeadline(t) 432 | return c.lastErr 433 | } 434 | 435 | func (c *tcpConn) SetWriteDeadline(t time.Time) error { 436 | c.lastErr = c.tcp.SetWriteDeadline(t) 437 | return c.lastErr 438 | } 439 | 440 | func (c *tcpConn) LastErr() error { 441 | return c.lastErr 442 | } 443 | 444 | func (c *tcpConn) GetLastUsed() time.Time { 445 | return c.lastUsed 446 | } 447 | 448 | func (c *tcpConn) SetLastUsed(t time.Time) { 449 | c.lastUsed = t 450 | } 451 | 452 | type maxValueCollector struct { 453 | values []time.Duration 454 | } 455 | 456 | func (m *maxValueCollector) Add(value time.Duration) { 457 | m.values = append(m.values, value) 458 | } 459 | 460 | func (m *maxValueCollector) Max() time.Duration { 461 | var max time.Duration 462 | for _, v := range m.values { 463 | if v > max { 464 | max = v 465 | } 466 | } 467 | return max 468 | } 469 | 470 | // split returns VALUE section and further (possibly partial) 471 | // data section if \r\n delimiter is present, otherwise returns nil 472 | func split(data []byte) [][]byte { 473 | parts := bytes.SplitAfterN(data, sep, 2) 474 | if len(parts) != 2 { 475 | return nil 476 | } 477 | return parts 478 | } 479 | 480 | // parseValueResp parses VALUE section of the response and returns 481 | // key, length of data section and whether an error occurred 482 | func parseValueResp(raw string) (string, int, error) { 483 | value := strings.TrimSuffix(raw, "\r\n") 484 | terms := strings.Split(value, " ") 485 | if terms == nil || strings.Compare(terms[0], "VALUE") == -1 || len(terms) != 4 { 486 | return "", 0, fmt.Errorf("incorrect value response %v", terms) 487 | } 488 | length, err := strconv.Atoi(terms[3]) 489 | if err != nil { 490 | return "", 0, fmt.Errorf("error parse data len field: %v", err) 491 | } 492 | return terms[1], length, nil 493 | } 494 | -------------------------------------------------------------------------------- /memcached_test.go: -------------------------------------------------------------------------------- 1 | package memcached 2 | 3 | import ( 4 | "context" 5 | "fmt" 6 | "math/rand" 7 | "net" 8 | "strconv" 9 | "strings" 10 | "testing" 11 | "time" 12 | ) 13 | 14 | func NewTestClient(addr string, maxIdleConn int, p connProvider) *Client { 15 | return &Client{ 16 | addr: addr, 17 | pool: newStack(maxIdleConn), 18 | provider: p, 19 | idleTimeout: defaultIdleConnectionTimeout, 20 | timeNow: time.Now, 21 | profiler: &nopProfiler{}, 22 | } 23 | } 24 | 25 | func TestGet(t *testing.T) { 26 | type testCase struct { 27 | name string 28 | key string 29 | data string 30 | dataChunks []string 31 | } 32 | 33 | testConn := mockConn{} 34 | provider := mockConnProvider{ 35 | conn: &testConn, 36 | } 37 | client := NewTestClient("localhost", 1, &provider) 38 | 39 | tests := []testCase{ 40 | { 41 | name: "unified", 42 | key: "key/3", 43 | data: "data/3", dataChunks: []string{"VALUE key/3 0 " + strconv.Itoa(len("data/3")) + "\r\ndata/3\r\nEND\r\n"}, 44 | }, 45 | { 46 | name: "partially chunked data", 47 | key: "key/1", 48 | data: "data/1", 49 | dataChunks: []string{"VALUE ", "key/1", " 0", " ", strconv.Itoa(len("data/1")), "\r\n", "data/1", "\r", "\n", "END\r", "\n"}, 50 | }, 51 | { 52 | name: "fully chunked data", 53 | key: "key/2", 54 | data: "data/2", 55 | dataChunks: []string{"V", "A", "L", "U", "E", ",", " ", "k", "e", "y", "/", "2", " ", "0", " ", strconv.Itoa(len("data/2")), "\r", "\n", "d", "a", "t", "a", "/", "2", "\r", "\n", "E", "N", "D", "\r", "\n"}, 56 | }, 57 | { 58 | name: "empty", 59 | key: "key/4", 60 | data: "", 61 | dataChunks: []string{"VALUE key/4 0 " + strconv.Itoa(len("")) + "\r\n\r\nEND\r\n"}, 62 | }, 63 | { 64 | name: "double-crlf", 65 | key: "key/5", 66 | data: "data/5\r\n", 67 | dataChunks: []string{"VALUE key/5 0 " + strconv.Itoa(len("data/5\r\n")) + "\r\ndata/5\r\n\r\nEND\r\n"}, 68 | }, 69 | { 70 | name: "extracr", 71 | key: "key/6", 72 | data: "data/6\r", 73 | dataChunks: []string{"VALUE key/6 0 " + strconv.Itoa(len("data/6\r")) + "\r\ndata/6\r\r\nEND\r\n"}, 74 | }, 75 | { 76 | name: "extralf", 77 | key: "key/7", 78 | data: "data/7\n", 79 | dataChunks: []string{"VALUE key/7 0 " + strconv.Itoa(len("data/7\n")) + "\r\ndata/7\n\r\nEND\r\n"}, 80 | }, 81 | { 82 | name: "crlf in the middle of the data", 83 | key: "key/8", 84 | data: "da\r\nta/8\n", 85 | dataChunks: []string{"VALUE key/7 0 " + strconv.Itoa(len("da\r\nta/8\n")) + "\r\nda\r\nta/8\n\r\nEND\r\n"}, 86 | }, 87 | { 88 | name: "ENDcrlf in the middle of the data", 89 | key: "key/9", 90 | data: "daEND\r\nta/9\n", 91 | dataChunks: []string{"VALUE key/9 0 " + strconv.Itoa(len("daEND\r\nta/9\n")) + "\r\ndaEND\r\nta/9\n\r\nEND\r\n"}, 92 | }, 93 | } 94 | 95 | ctx, cancel := context.WithCancel(context.Background()) 96 | defer cancel() 97 | 98 | for _, test := range tests { 99 | t.Run(test.name, func(t *testing.T) { 100 | defer testConn.Reset() 101 | testConn.dataChunks = test.dataChunks 102 | rsp, err := client.Get(ctx, test.key) 103 | if err != nil { 104 | t.Fatal("get error:", err) 105 | } 106 | if rsp.Data == nil || string(rsp.Data) != test.data { 107 | t.Fatalf("expected data field %v but got %v", test.data, string(rsp.Data)) 108 | } 109 | }) 110 | } 111 | } 112 | 113 | func TestGetJunkData(t *testing.T) { 114 | type testCase struct { 115 | name string 116 | key string 117 | length int 118 | } 119 | 120 | testConn := mockConnDataGenerator{} 121 | provider := mockConnProvider{ 122 | conn: &testConn, 123 | } 124 | client := NewTestClient("localhost", 1, &provider) 125 | 126 | tests := []testCase{ 127 | { 128 | name: "generator", 129 | key: "key/1", 130 | length: internalBuffSize * 2, 131 | }, 132 | } 133 | 134 | ctx, cancel := context.WithCancel(context.Background()) 135 | defer cancel() 136 | 137 | for _, test := range tests { 138 | t.Run(test.name, func(t *testing.T) { 139 | defer testConn.Reset() 140 | testConn.JunkData = make([]byte, test.length) 141 | _, err := client.Get(ctx, test.key) 142 | if err == nil { 143 | t.Fatal("expected buffer overflow error") 144 | } 145 | }) 146 | } 147 | } 148 | 149 | func TestGetTimeout(t *testing.T) { 150 | type testCase struct { 151 | name string 152 | key string 153 | dataChunks []string 154 | invoked bool 155 | } 156 | 157 | testConn := mockConn{} 158 | provider := mockConnProvider{ 159 | conn: &testConn, 160 | } 161 | client := NewTestClient("localhost", 1, &provider) 162 | 163 | tests := []testCase{ 164 | { 165 | name: "timeout", 166 | key: "test/1", 167 | dataChunks: []string{"VALUE key/1 0 " + strconv.Itoa(len("data/1")) + "\r\ndata/1\r\nEND\r\n"}, 168 | invoked: true, 169 | }, 170 | } 171 | 172 | deadline := time.Now().Add(1 * time.Second) 173 | ctx, cancel := context.WithDeadline(context.Background(), deadline) 174 | defer cancel() 175 | 176 | for _, test := range tests { 177 | t.Run(test.name, func(t *testing.T) { 178 | defer testConn.Reset() 179 | testConn.dataChunks = test.dataChunks 180 | _, err := client.Get(ctx, test.key) 181 | if err != nil { 182 | t.Fatalf("get error: %v", err) 183 | } 184 | if testConn.readDeadline != test.invoked { 185 | t.Fatalf("expected get timeout") 186 | } 187 | if testConn.deadline != deadline { 188 | t.Fatalf("expected deadline %v but got %v", deadline, testConn.deadline) 189 | } 190 | }) 191 | } 192 | } 193 | 194 | func TestGetConcurrencyRace(t *testing.T) { 195 | type testCase struct { 196 | name string 197 | key string 198 | data string 199 | dataChunks []string 200 | } 201 | 202 | provider := mockMultiConnProvider{} 203 | client := NewTestClient("localhost", 1, &provider) 204 | 205 | test := testCase{ 206 | name: "fully chunked data", 207 | key: "key/2", 208 | data: "data/2", 209 | dataChunks: []string{"V", "A", "L", "U", "E", " ", "k", "e", "y", "/", "2", " ", "0", " ", strconv.Itoa(len("data/2")), "\r", "\n", "d", "a", "t", "a", "/", "2", "\r", "\n", "E", "N", "D", "\r", "\n"}, 210 | } 211 | 212 | provider.dataChunks = test.dataChunks 213 | 214 | for i := 0; i < 8000; i++ { // 8128 is a limit on simultaneously alive goroutines with -race 215 | test2 := test 216 | t.Run(fmt.Sprintf("test_%v", i+1), func(t *testing.T) { 217 | t.Parallel() 218 | rsp, err := client.Get(context.Background(), test2.key) 219 | if err != nil { 220 | t.Fatal("get error:", err) 221 | } 222 | if rsp.Data == nil || string(rsp.Data) != test2.data { 223 | t.Fatalf("expected data field %v but got %v", test.data, string(rsp.Data)) 224 | } 225 | }) 226 | } 227 | } 228 | 229 | func TestSet(t *testing.T) { 230 | type testCase struct { 231 | name string 232 | key string 233 | data []byte 234 | exp int 235 | expected string 236 | } 237 | 238 | testConn := mockConn{writeMode: true} 239 | provider := mockConnProvider{ 240 | conn: &testConn, 241 | } 242 | client := NewTestClient("localhost", 1, &provider) 243 | 244 | tests := []testCase{ 245 | { 246 | name: "key 3", 247 | key: "key/3", 248 | data: []byte("data/3"), 249 | exp: 3, 250 | expected: "set key/3 0 3 6\r\ndata/3\r\n", 251 | }, 252 | { 253 | name: "key 1", 254 | key: "key/1", 255 | data: []byte("data/1"), 256 | exp: 4, 257 | expected: "set key/1 0 4 6\r\ndata/1\r\n", 258 | }, 259 | { 260 | name: "key 2", 261 | key: "key/2", 262 | data: []byte("data/2"), 263 | exp: 5, 264 | expected: "set key/2 0 5 6\r\ndata/2\r\n", 265 | }, 266 | } 267 | 268 | ctx, cancel := context.WithCancel(context.Background()) 269 | defer cancel() 270 | 271 | for _, test := range tests { 272 | t.Run(test.name, func(t *testing.T) { 273 | defer func() { 274 | testConn.dataChunks = nil 275 | testConn.Reset() 276 | }() 277 | if err := client.Set(ctx, test.key, test.data, test.exp); err != nil { 278 | t.Fatal("set error:", err) 279 | } 280 | data := strings.Join(testConn.dataChunks, "") 281 | if data != test.expected { 282 | t.Fatalf("expected: [%v] but got: [%v]\n", prettify(test.expected), prettify(data)) 283 | } 284 | }) 285 | } 286 | } 287 | 288 | func TestIdleConnTimeout(t *testing.T) { 289 | type testCase struct { 290 | name string 291 | time time.Time 292 | expectNewConnection bool 293 | } 294 | 295 | testConn := mockConn{writeMode: true} 296 | provider := mockConnProvider{ 297 | conn: &testConn, 298 | } 299 | client := NewTestClient("localhost", 1, &provider) 300 | client.idleTimeout = 1 * time.Second 301 | 302 | t1, _ := time.Parse(time.RFC3339, "2006-01-02T15:04:05Z") 303 | t2, _ := time.Parse(time.RFC3339, "2006-01-02T15:04:07Z") 304 | t3, _ := time.Parse(time.RFC3339, "2006-01-02T15:04:07Z") 305 | t4, _ := time.Parse(time.RFC3339, "2006-01-02T15:04:09Z") 306 | t5, _ := time.Parse(time.RFC3339, "2006-01-02T15:04:11Z") 307 | 308 | tests := []testCase{ 309 | { 310 | name: "First new connection", 311 | time: t1, 312 | expectNewConnection: true, 313 | }, 314 | { 315 | name: "Second new connection", 316 | time: t2, 317 | expectNewConnection: true, 318 | }, 319 | { 320 | name: "Old connection", 321 | time: t3, 322 | expectNewConnection: false, 323 | }, 324 | { 325 | name: "Third new connection", 326 | time: t4, 327 | expectNewConnection: true, 328 | }, 329 | { 330 | name: "Fourth new connection", 331 | time: t5, 332 | expectNewConnection: true, 333 | }, 334 | } 335 | 336 | ctx, cancel := context.WithCancel(context.Background()) 337 | defer cancel() 338 | 339 | for _, test := range tests { 340 | client.timeNow = func() time.Time { 341 | return test.time 342 | } 343 | if err := client.Set(ctx, "key", []byte("data"), 5); err != nil { 344 | t.Fatal("set error:", err) 345 | } 346 | if provider.newConnCreated != test.expectNewConnection { 347 | t.Fatalf("expected new connection created: [%v] but got: [%v]\n", test.expectNewConnection, provider.newConnCreated) 348 | } 349 | testConn.Reset() 350 | provider.Reset() 351 | testConn.dataChunks = nil 352 | } 353 | } 354 | 355 | func TestConnReused(t *testing.T) { 356 | const maxIdleConn = 100 357 | 358 | testConn := mockConn{writeMode: true} 359 | provider := mockConnProvider{ 360 | conn: &testConn, 361 | } 362 | client := NewTestClient("localhost", maxIdleConn, &provider) 363 | client.idleTimeout = 500 * time.Millisecond 364 | 365 | ctx, cancel := context.WithCancel(context.Background()) 366 | defer cancel() 367 | 368 | currentTime, _ := time.Parse(time.RFC3339, "2006-01-02T15:04:05Z") 369 | var connTimes []time.Time 370 | 371 | for i := 0; i < 100; i++ { 372 | msec := rand.Intn(999) 373 | newTime := currentTime.Add(time.Duration(msec) * time.Millisecond) 374 | client.timeNow = func() time.Time { 375 | return newTime 376 | } 377 | 378 | expectNewConnection, index := willUseNewConnection(connTimes, newTime, client.idleTimeout) 379 | 380 | if expectNewConnection { 381 | connTimes = nil 382 | connTimes = append(connTimes, newTime) 383 | } else { 384 | connTimes[index] = newTime 385 | connTimes = connTimes[index:] 386 | } 387 | 388 | if err := client.Set(ctx, "key", []byte("data"), 5); err != nil { 389 | t.Fatal("set error:", err) 390 | } 391 | 392 | if provider.newConnCreated != expectNewConnection { 393 | t.Fatalf("expected new connection created: [%v] but got: [%v]\n", expectNewConnection, provider.newConnCreated) 394 | } 395 | 396 | testConn.Reset() 397 | provider.Reset() 398 | testConn.dataChunks = nil 399 | 400 | currentTime = newTime 401 | } 402 | } 403 | 404 | func willUseNewConnection(connTimes []time.Time, newTime time.Time, timeout time.Duration) (bool, int) { 405 | for i, t := range connTimes { 406 | if newTime.Sub(t) < timeout { 407 | return false, i 408 | } 409 | } 410 | return true, -1 411 | } 412 | 413 | func prettify(str string) string { 414 | return strings.ReplaceAll(str, "\r\n", "\\r\\n") 415 | } 416 | 417 | type mockConnProvider struct { 418 | conn memConn 419 | newConnCreated bool 420 | } 421 | 422 | func (p *mockConnProvider) NewConn(address string) (memConn, error) { 423 | p.newConnCreated = true 424 | return p.conn, nil 425 | } 426 | 427 | func (p *mockConnProvider) Reset() { 428 | p.newConnCreated = false 429 | } 430 | 431 | type mockMultiConnProvider struct { 432 | dataChunks []string 433 | } 434 | 435 | func (p *mockMultiConnProvider) NewConn(address string) (memConn, error) { 436 | conn := mockConn{dataChunks: p.dataChunks} 437 | return &conn, nil 438 | } 439 | 440 | type mockConn struct { 441 | dataChunks []string 442 | counter int 443 | readDeadline bool 444 | writeDeadline bool 445 | deadline time.Time 446 | writeMode bool 447 | lastUsed time.Time 448 | } 449 | 450 | func (c *mockConn) Reset() { 451 | c.counter = 0 452 | c.readDeadline = false 453 | c.writeDeadline = false 454 | c.deadline = time.Time{} 455 | } 456 | 457 | func (c *mockConn) Read(b []byte) (n int, err error) { 458 | var bytesCopied int 459 | if c.writeMode { 460 | str := "STORED\r\n" 461 | raw := []byte(str) 462 | if len(raw) > len(b) { 463 | raw = raw[:len(b)] 464 | } 465 | bytesCopied = copy(b, raw[c.counter:]) 466 | c.counter += bytesCopied 467 | } else { 468 | if c.counter == len(c.dataChunks) { 469 | c.Reset() 470 | } 471 | str := c.dataChunks[c.counter] 472 | c.counter++ 473 | raw := []byte(str) 474 | if len(raw) > len(b) { 475 | raw = raw[:len(b)] 476 | } 477 | bytesCopied = copy(b, raw) 478 | } 479 | return bytesCopied, nil 480 | } 481 | 482 | func (c *mockConn) Write(b []byte) (n int, err error) { 483 | var bytesCopied int 484 | if c.writeMode { 485 | buf := make([]byte, len(b)) 486 | bytesCopied = copy(buf, b) 487 | c.dataChunks = append(c.dataChunks, string(buf)) 488 | } 489 | return bytesCopied, nil 490 | } 491 | 492 | func (c *mockConn) Close() error { 493 | return nil 494 | } 495 | 496 | func (c *mockConn) LocalAddr() net.Addr { 497 | return nil 498 | } 499 | 500 | func (c *mockConn) RemoteAddr() net.Addr { 501 | return nil 502 | } 503 | 504 | func (c *mockConn) SetDeadline(t time.Time) error { 505 | c.readDeadline = true 506 | c.writeDeadline = true 507 | c.deadline = t 508 | return nil 509 | } 510 | 511 | func (c *mockConn) SetReadDeadline(t time.Time) error { 512 | c.readDeadline = true 513 | c.deadline = t 514 | return nil 515 | } 516 | 517 | func (c *mockConn) SetWriteDeadline(t time.Time) error { 518 | c.writeDeadline = true 519 | c.deadline = t 520 | return nil 521 | } 522 | 523 | func (c *mockConn) LastErr() error { 524 | return nil 525 | } 526 | 527 | func (c *mockConn) GetLastUsed() time.Time { 528 | return c.lastUsed 529 | } 530 | 531 | func (c *mockConn) SetLastUsed(t time.Time) { 532 | c.lastUsed = t 533 | } 534 | 535 | type mockConnDataGenerator struct { 536 | JunkData []byte 537 | pointer int 538 | lastUsed time.Time 539 | } 540 | 541 | func (c *mockConnDataGenerator) Reset() { 542 | c.pointer = 0 543 | } 544 | 545 | func (c *mockConnDataGenerator) Read(b []byte) (n int, err error) { 546 | if c.JunkData == nil { 547 | return 0, fmt.Errorf("junk data not allocated") 548 | } 549 | if len(c.JunkData) == c.pointer { 550 | // return 0, io.EOF 551 | c.Reset() 552 | } 553 | var end int 554 | if len(b) < len(c.JunkData)-c.pointer { 555 | end = c.pointer + len(b) 556 | } else { 557 | end = len(c.JunkData) 558 | } 559 | bytesCopied := copy(b, c.JunkData[c.pointer:end]) 560 | c.pointer += bytesCopied 561 | return bytesCopied, nil 562 | } 563 | 564 | func (c *mockConnDataGenerator) Write(b []byte) (n int, err error) { 565 | return -1, nil 566 | } 567 | 568 | func (c *mockConnDataGenerator) Close() error { 569 | return nil 570 | } 571 | 572 | func (c *mockConnDataGenerator) LocalAddr() net.Addr { 573 | return nil 574 | } 575 | 576 | func (c *mockConnDataGenerator) RemoteAddr() net.Addr { 577 | return nil 578 | } 579 | 580 | func (c *mockConnDataGenerator) SetDeadline(t time.Time) error { 581 | return nil 582 | } 583 | 584 | func (c *mockConnDataGenerator) SetReadDeadline(t time.Time) error { 585 | return nil 586 | } 587 | 588 | func (c *mockConnDataGenerator) SetWriteDeadline(t time.Time) error { 589 | return nil 590 | } 591 | 592 | func (c *mockConnDataGenerator) LastErr() error { 593 | return nil 594 | } 595 | 596 | func (c *mockConnDataGenerator) GetLastUsed() time.Time { 597 | return c.lastUsed 598 | } 599 | 600 | func (c *mockConnDataGenerator) SetLastUsed(t time.Time) { 601 | c.lastUsed = t 602 | } 603 | --------------------------------------------------------------------------------