diff --git a/Gopkg.lock b/Gopkg.lock index f88447a5..1e11a33e 100644 --- a/Gopkg.lock +++ b/Gopkg.lock @@ -37,6 +37,12 @@ packages = ["."] revision = "855e8d98f1852d48dde521e0522408d1fe7e836a" +[[projects]] + branch = "master" + name = "github.com/hashicorp/go-immutable-radix" + packages = ["."] + revision = "59b67882ec612f43b9d4c4fd97cebd507be4b3ee" + [[projects]] branch = "master" name = "github.com/hashicorp/golang-lru" @@ -99,6 +105,6 @@ [solve-meta] analyzer-name = "dep" analyzer-version = 1 - inputs-digest = "9946fe30a0b048dbe5b8a10b28e8bffd7ec6dac56380db345cbd868862fe7f08" + inputs-digest = "01e41ba3f8bb51f155fbca27baa1f94d4bc8e8d21d094b531941d78e01945fe2" solver-name = "gps-cdcl" solver-version = 1 diff --git a/Gopkg.toml b/Gopkg.toml index d47307de..232eaafa 100644 --- a/Gopkg.toml +++ b/Gopkg.toml @@ -37,3 +37,7 @@ [[constraint]] branch = "master" name = "golang.org/x/crypto" + +[[constraint]] + branch = "master" + name = "github.com/hashicorp/go-immutable-radix" diff --git a/dnscrypt-proxy/common.go b/dnscrypt-proxy/common.go index 81277ebd..f67a3e6e 100644 --- a/dnscrypt-proxy/common.go +++ b/dnscrypt-proxy/common.go @@ -72,3 +72,11 @@ func Max(a, b int) int { } return b } + +func StringReverse(s string) string { + r := []rune(s) + for i, j := 0, len(r)-1; i < len(r)/2; i, j = i+1, j-1 { + r[i], r[j] = r[j], r[i] + } + return string(r) +} diff --git a/dnscrypt-proxy/config.go b/dnscrypt-proxy/config.go index b78058a0..ca9031c4 100644 --- a/dnscrypt-proxy/config.go +++ b/dnscrypt-proxy/config.go @@ -25,6 +25,7 @@ type Config struct { CacheMinTTL uint32 `toml:"cache_min_ttl"` CacheMaxTTL uint32 `toml:"cache_max_ttl"` QueryLog QueryLogConfig `toml:"query_log"` + BlockName BlockNameConfig `toml:"block_name"` ServersConfig map[string]ServerConfig `toml:"servers"` SourcesConfig map[string]SourceConfig `toml:"sources"` } @@ -64,6 +65,10 @@ type QueryLogConfig struct { Format string } +type BlockNameConfig struct { + File string +} + func ConfigLoad(proxy *Proxy, config_file string) error { configFile := flag.String("config", "dnscrypt-proxy.toml", "path to the configuration file") flag.Parse() @@ -98,6 +103,7 @@ func ConfigLoad(proxy *Proxy, config_file string) error { } proxy.queryLogFile = config.QueryLog.File proxy.queryLogFormat = config.QueryLog.Format + proxy.blockNameFile = config.BlockName.File if len(config.ServerNames) == 0 { for serverName := range config.ServersConfig { config.ServerNames = append(config.ServerNames, serverName) diff --git a/dnscrypt-proxy/dnscrypt-proxy.toml b/dnscrypt-proxy/dnscrypt-proxy.toml index c161cb41..a471be85 100644 --- a/dnscrypt-proxy/dnscrypt-proxy.toml +++ b/dnscrypt-proxy/dnscrypt-proxy.toml @@ -54,12 +54,19 @@ block_ipv6 = false [query_log] ### Full path to the query log file -file = "/tmp/query.log" +# file = "/tmp/query.log" ### Query log format (currently supported: tsv and ltsv) format = "tsv" +############## Pattern-based blocking ############## + +[block_name] +## Full path to the file of blocking rules +# file = "/tmp/mybase.txt" + + ############## DNS Cache ############## ## Enable a basic DNS cache to reduce outgoing traffic diff --git a/dnscrypt-proxy/dnsutils.go b/dnscrypt-proxy/dnsutils.go index 96bad9c9..bf915875 100644 --- a/dnscrypt-proxy/dnsutils.go +++ b/dnscrypt-proxy/dnsutils.go @@ -1,6 +1,7 @@ package main import ( + "strings" "time" "github.com/miekg/dns" @@ -29,6 +30,15 @@ func EmptyResponseFromMessage(srcMsg *dns.Msg) (*dns.Msg, error) { return dstMsg, nil } +func RefusedResponseFromMessage(srcMsg *dns.Msg) (*dns.Msg, error) { + dstMsg, err := EmptyResponseFromMessage(srcMsg) + if err != nil { + return dstMsg, err + } + dstMsg.Rcode = dns.RcodeRefused + return dstMsg, nil +} + func HasTCFlag(packet []byte) bool { return packet[2]&2 == 2 } @@ -41,6 +51,13 @@ func NormalizeName(name *[]byte) { } } +func StripTrailingDot(str string) string { + if strings.HasSuffix(str, ".") { + str = str[:len(str)-1] + } + return str +} + func getMinTTL(msg *dns.Msg, minTTL uint32, maxTTL uint32, negCacheMinTTL uint32) time.Duration { if msg.Rcode != dns.RcodeSuccess || len(msg.Answer) <= 0 { return time.Duration(negCacheMinTTL) * time.Second diff --git a/dnscrypt-proxy/main.go b/dnscrypt-proxy/main.go index 0e852b78..f70d8f14 100644 --- a/dnscrypt-proxy/main.go +++ b/dnscrypt-proxy/main.go @@ -30,6 +30,7 @@ type Proxy struct { cacheMaxTTL uint32 queryLogFile string queryLogFormat string + blockNameFile string pluginsGlobals PluginsGlobals } @@ -194,6 +195,9 @@ func (proxy *Proxy) processIncomingQuery(serverInfo *ServerInfo, clientProto str return } } + if pluginsState.action == PluginsActionDrop { + return + } } if len(response) == 0 { encryptedQuery, clientNonce, err := proxy.Encrypt(serverInfo, query, serverProto) diff --git a/dnscrypt-proxy/plugin_block_name.go b/dnscrypt-proxy/plugin_block_name.go new file mode 100644 index 00000000..ce392a38 --- /dev/null +++ b/dnscrypt-proxy/plugin_block_name.go @@ -0,0 +1,129 @@ +package main + +import ( + "io/ioutil" + "strings" + + "github.com/hashicorp/go-immutable-radix" + "github.com/jedisct1/dlog" + "github.com/miekg/dns" +) + +type PluginBlockType int + +const ( + PluginBlockTypeNone = iota + PluginBlockTypePrefix + PluginBlockTypeSuffix + PluginBlockTypeSubstring + PluginBlockTypePattern +) + +type PluginBlockName struct { + blockedPrefixes *iradix.Tree + blockedSuffixes *iradix.Tree + blockedSubstrings []string + blockedPatterns []string +} + +func (plugin *PluginBlockName) Name() string { + return "block_name" +} + +func (plugin *PluginBlockName) Description() string { + return "Block DNS queries matching name patterns" +} + +func (plugin *PluginBlockName) Init(proxy *Proxy) error { + dlog.Noticef("Loading the set of blocking rules from [%s]", proxy.blockNameFile) + bin, err := ioutil.ReadFile(proxy.blockNameFile) + if err != nil { + return err + } + plugin.blockedPrefixes = iradix.New() + plugin.blockedSuffixes = iradix.New() + for lineNo, line := range strings.Split(string(bin), "\n") { + line = strings.Trim(line, " \t\r\n") + if len(line) == 0 || strings.HasPrefix(line, "#") { + continue + } + leadingStar := strings.HasPrefix(line, "*") + trailingStar := strings.HasSuffix(line, "*") + blockType := PluginBlockTypeNone + if leadingStar && trailingStar { + blockType = PluginBlockTypeSubstring + if len(line) < 3 { + dlog.Errorf("Syntax error in block rules at line %d", lineNo) + continue + } + line = line[1 : len(line)-1] + } else if trailingStar { + blockType = PluginBlockTypePrefix + if len(line) < 2 { + dlog.Errorf("Syntax error in block rules at line %d", lineNo) + continue + } + line = line[:len(line)-1] + } else { + blockType = PluginBlockTypeSuffix + if leadingStar { + line = line[1:] + } + if strings.HasPrefix(line, ".") { + line = line[1:] + } + } + if len(line) == 0 { + dlog.Errorf("Syntax error in block rule at line %d", lineNo) + continue + } + line = strings.ToLower(line) + switch blockType { + case PluginBlockTypeSubstring: + plugin.blockedSubstrings = append(plugin.blockedSubstrings, line) + case PluginBlockTypePrefix: + plugin.blockedPrefixes, _, _ = plugin.blockedPrefixes.Insert([]byte(line), 0) + case PluginBlockTypeSuffix: + plugin.blockedSuffixes, _, _ = plugin.blockedSuffixes.Insert([]byte(StringReverse(line)), 0) + default: + dlog.Fatal("Unexpected block type") + } + } + return nil +} + +func (plugin *PluginBlockName) Drop() error { + return nil +} + +func (plugin *PluginBlockName) Reload() error { + return nil +} + +func (plugin *PluginBlockName) Eval(pluginsState *PluginsState, msg *dns.Msg) error { + questions := msg.Question + if len(questions) != 1 { + return nil + } + question := strings.ToLower(StripTrailingDot(questions[0].Name)) + revQuestion := StringReverse(question) + match, _, found := plugin.blockedSuffixes.Root().LongestPrefix([]byte(revQuestion)) + if found { + if len(match) == len(question) || question[len(match)] == '.' { + pluginsState.action = PluginsActionReject + return nil + } + } + _, _, found = plugin.blockedPrefixes.Root().LongestPrefix([]byte(question)) + if found { + pluginsState.action = PluginsActionReject + return nil + } + for _, substring := range plugin.blockedSubstrings { + if strings.Contains(substring, question) { + pluginsState.action = PluginsActionReject + return nil + } + } + return nil +} diff --git a/dnscrypt-proxy/plugins.go b/dnscrypt-proxy/plugins.go index 7d4cb8ba..92f7d71d 100644 --- a/dnscrypt-proxy/plugins.go +++ b/dnscrypt-proxy/plugins.go @@ -1,9 +1,11 @@ package main import ( + "errors" "net" "sync" + "github.com/jedisct1/dlog" "github.com/miekg/dns" ) @@ -43,6 +45,9 @@ func InitPluginsGlobals(pluginsGlobals *PluginsGlobals, proxy *Proxy) error { if len(proxy.queryLogFile) != 0 { *queryPlugins = append(*queryPlugins, Plugin(new(PluginQueryLog))) } + if len(proxy.blockNameFile) != 0 { + *queryPlugins = append(*queryPlugins, Plugin(new(PluginBlockName))) + } if proxy.pluginBlockIPv6 { *queryPlugins = append(*queryPlugins, Plugin(new(PluginBlockIPv6))) } @@ -103,6 +108,9 @@ func (pluginsState *PluginsState) ApplyQueryPlugins(pluginsGlobals *PluginsGloba if err := msg.Unpack(packet); err != nil { return packet, err } + if len(msg.Question) > 1 { + return packet, errors.New("Unexpected number of questions") + } pluginsGlobals.RLock() for _, plugin := range *pluginsGlobals.queryPlugins { if ret := plugin.Eval(pluginsState, &msg); ret != nil { @@ -110,6 +118,13 @@ func (pluginsState *PluginsState) ApplyQueryPlugins(pluginsGlobals *PluginsGloba pluginsState.action = PluginsActionDrop return packet, ret } + if pluginsState.action == PluginsActionReject { + synth, err := RefusedResponseFromMessage(&msg) + if err != nil { + return nil, err + } + pluginsState.synthResponse = synth + } if pluginsState.action != PluginsActionForward { break } @@ -138,6 +153,14 @@ func (pluginsState *PluginsState) ApplyResponsePlugins(pluginsGlobals *PluginsGl pluginsState.action = PluginsActionDrop return packet, ret } + if pluginsState.action == PluginsActionReject { + synth, err := RefusedResponseFromMessage(&msg) + if err != nil { + return nil, err + } + dlog.Infof("Blocking [%s]", synth.Question[0].Name) + pluginsState.synthResponse = synth + } if pluginsState.action != PluginsActionForward { break } diff --git a/dnscrypt-proxy/sources.go b/dnscrypt-proxy/sources.go index 1ba8c541..fd2c6c1a 100644 --- a/dnscrypt-proxy/sources.go +++ b/dnscrypt-proxy/sources.go @@ -131,14 +131,14 @@ func (source *Source) Parse() ([]RegisteredServer, error) { if err != nil { return registeredServers, nil } - for line, record := range records { + for lineNo, record := range records { if len(record) == 0 { continue } if len(record) < 14 { - return registeredServers, fmt.Errorf("Parse error at line %d", line) + return registeredServers, fmt.Errorf("Parse error at line %d", lineNo) } - if line == 0 { + if lineNo == 0 { continue } name := record[0] diff --git a/vendor/github.com/hashicorp/go-immutable-radix/.gitignore b/vendor/github.com/hashicorp/go-immutable-radix/.gitignore new file mode 100644 index 00000000..daf913b1 --- /dev/null +++ b/vendor/github.com/hashicorp/go-immutable-radix/.gitignore @@ -0,0 +1,24 @@ +# Compiled Object files, Static and Dynamic libs (Shared Objects) +*.o +*.a +*.so + +# Folders +_obj +_test + +# Architecture specific extensions/prefixes +*.[568vq] +[568vq].out + +*.cgo1.go +*.cgo2.c +_cgo_defun.c +_cgo_gotypes.go +_cgo_export.* + +_testmain.go + +*.exe +*.test +*.prof diff --git a/vendor/github.com/hashicorp/go-immutable-radix/.travis.yml b/vendor/github.com/hashicorp/go-immutable-radix/.travis.yml new file mode 100644 index 00000000..1a0bbea6 --- /dev/null +++ b/vendor/github.com/hashicorp/go-immutable-radix/.travis.yml @@ -0,0 +1,3 @@ +language: go +go: + - tip diff --git a/vendor/github.com/hashicorp/go-immutable-radix/LICENSE b/vendor/github.com/hashicorp/go-immutable-radix/LICENSE new file mode 100644 index 00000000..e87a115e --- /dev/null +++ b/vendor/github.com/hashicorp/go-immutable-radix/LICENSE @@ -0,0 +1,363 @@ +Mozilla Public License, version 2.0 + +1. Definitions + +1.1. "Contributor" + + means each individual or legal entity that creates, contributes to the + creation of, or owns Covered Software. + +1.2. "Contributor Version" + + means the combination of the Contributions of others (if any) used by a + Contributor and that particular Contributor's Contribution. + +1.3. "Contribution" + + means Covered Software of a particular Contributor. + +1.4. "Covered Software" + + means Source Code Form to which the initial Contributor has attached the + notice in Exhibit A, the Executable Form of such Source Code Form, and + Modifications of such Source Code Form, in each case including portions + thereof. + +1.5. "Incompatible With Secondary Licenses" + means + + a. that the initial Contributor has attached the notice described in + Exhibit B to the Covered Software; or + + b. that the Covered Software was made available under the terms of + version 1.1 or earlier of the License, but not also under the terms of + a Secondary License. + +1.6. "Executable Form" + + means any form of the work other than Source Code Form. + +1.7. "Larger Work" + + means a work that combines Covered Software with other material, in a + separate file or files, that is not Covered Software. + +1.8. "License" + + means this document. + +1.9. "Licensable" + + means having the right to grant, to the maximum extent possible, whether + at the time of the initial grant or subsequently, any and all of the + rights conveyed by this License. + +1.10. "Modifications" + + means any of the following: + + a. any file in Source Code Form that results from an addition to, + deletion from, or modification of the contents of Covered Software; or + + b. any new file in Source Code Form that contains any Covered Software. + +1.11. "Patent Claims" of a Contributor + + means any patent claim(s), including without limitation, method, + process, and apparatus claims, in any patent Licensable by such + Contributor that would be infringed, but for the grant of the License, + by the making, using, selling, offering for sale, having made, import, + or transfer of either its Contributions or its Contributor Version. + +1.12. "Secondary License" + + means either the GNU General Public License, Version 2.0, the GNU Lesser + General Public License, Version 2.1, the GNU Affero General Public + License, Version 3.0, or any later versions of those licenses. + +1.13. "Source Code Form" + + means the form of the work preferred for making modifications. + +1.14. "You" (or "Your") + + means an individual or a legal entity exercising rights under this + License. For legal entities, "You" includes any entity that controls, is + controlled by, or is under common control with You. For purposes of this + definition, "control" means (a) the power, direct or indirect, to cause + the direction or management of such entity, whether by contract or + otherwise, or (b) ownership of more than fifty percent (50%) of the + outstanding shares or beneficial ownership of such entity. + + +2. License Grants and Conditions + +2.1. Grants + + Each Contributor hereby grants You a world-wide, royalty-free, + non-exclusive license: + + a. under intellectual property rights (other than patent or trademark) + Licensable by such Contributor to use, reproduce, make available, + modify, display, perform, distribute, and otherwise exploit its + Contributions, either on an unmodified basis, with Modifications, or + as part of a Larger Work; and + + b. under Patent Claims of such Contributor to make, use, sell, offer for + sale, have made, import, and otherwise transfer either its + Contributions or its Contributor Version. + +2.2. Effective Date + + The licenses granted in Section 2.1 with respect to any Contribution + become effective for each Contribution on the date the Contributor first + distributes such Contribution. + +2.3. Limitations on Grant Scope + + The licenses granted in this Section 2 are the only rights granted under + this License. No additional rights or licenses will be implied from the + distribution or licensing of Covered Software under this License. + Notwithstanding Section 2.1(b) above, no patent license is granted by a + Contributor: + + a. for any code that a Contributor has removed from Covered Software; or + + b. for infringements caused by: (i) Your and any other third party's + modifications of Covered Software, or (ii) the combination of its + Contributions with other software (except as part of its Contributor + Version); or + + c. under Patent Claims infringed by Covered Software in the absence of + its Contributions. + + This License does not grant any rights in the trademarks, service marks, + or logos of any Contributor (except as may be necessary to comply with + the notice requirements in Section 3.4). + +2.4. Subsequent Licenses + + No Contributor makes additional grants as a result of Your choice to + distribute the Covered Software under a subsequent version of this + License (see Section 10.2) or under the terms of a Secondary License (if + permitted under the terms of Section 3.3). + +2.5. Representation + + Each Contributor represents that the Contributor believes its + Contributions are its original creation(s) or it has sufficient rights to + grant the rights to its Contributions conveyed by this License. + +2.6. Fair Use + + This License is not intended to limit any rights You have under + applicable copyright doctrines of fair use, fair dealing, or other + equivalents. + +2.7. Conditions + + Sections 3.1, 3.2, 3.3, and 3.4 are conditions of the licenses granted in + Section 2.1. + + +3. Responsibilities + +3.1. Distribution of Source Form + + All distribution of Covered Software in Source Code Form, including any + Modifications that You create or to which You contribute, must be under + the terms of this License. You must inform recipients that the Source + Code Form of the Covered Software is governed by the terms of this + License, and how they can obtain a copy of this License. You may not + attempt to alter or restrict the recipients' rights in the Source Code + Form. + +3.2. Distribution of Executable Form + + If You distribute Covered Software in Executable Form then: + + a. such Covered Software must also be made available in Source Code Form, + as described in Section 3.1, and You must inform recipients of the + Executable Form how they can obtain a copy of such Source Code Form by + reasonable means in a timely manner, at a charge no more than the cost + of distribution to the recipient; and + + b. You may distribute such Executable Form under the terms of this + License, or sublicense it under different terms, provided that the + license for the Executable Form does not attempt to limit or alter the + recipients' rights in the Source Code Form under this License. + +3.3. Distribution of a Larger Work + + You may create and distribute a Larger Work under terms of Your choice, + provided that You also comply with the requirements of this License for + the Covered Software. If the Larger Work is a combination of Covered + Software with a work governed by one or more Secondary Licenses, and the + Covered Software is not Incompatible With Secondary Licenses, this + License permits You to additionally distribute such Covered Software + under the terms of such Secondary License(s), so that the recipient of + the Larger Work may, at their option, further distribute the Covered + Software under the terms of either this License or such Secondary + License(s). + +3.4. Notices + + You may not remove or alter the substance of any license notices + (including copyright notices, patent notices, disclaimers of warranty, or + limitations of liability) contained within the Source Code Form of the + Covered Software, except that You may alter any license notices to the + extent required to remedy known factual inaccuracies. + +3.5. Application of Additional Terms + + You may choose to offer, and to charge a fee for, warranty, support, + indemnity or liability obligations to one or more recipients of Covered + Software. However, You may do so only on Your own behalf, and not on + behalf of any Contributor. You must make it absolutely clear that any + such warranty, support, indemnity, or liability obligation is offered by + You alone, and You hereby agree to indemnify every Contributor for any + liability incurred by such Contributor as a result of warranty, support, + indemnity or liability terms You offer. You may include additional + disclaimers of warranty and limitations of liability specific to any + jurisdiction. + +4. Inability to Comply Due to Statute or Regulation + + If it is impossible for You to comply with any of the terms of this License + with respect to some or all of the Covered Software due to statute, + judicial order, or regulation then You must: (a) comply with the terms of + this License to the maximum extent possible; and (b) describe the + limitations and the code they affect. Such description must be placed in a + text file included with all distributions of the Covered Software under + this License. Except to the extent prohibited by statute or regulation, + such description must be sufficiently detailed for a recipient of ordinary + skill to be able to understand it. + +5. Termination + +5.1. The rights granted under this License will terminate automatically if You + fail to comply with any of its terms. However, if You become compliant, + then the rights granted under this License from a particular Contributor + are reinstated (a) provisionally, unless and until such Contributor + explicitly and finally terminates Your grants, and (b) on an ongoing + basis, if such Contributor fails to notify You of the non-compliance by + some reasonable means prior to 60 days after You have come back into + compliance. Moreover, Your grants from a particular Contributor are + reinstated on an ongoing basis if such Contributor notifies You of the + non-compliance by some reasonable means, this is the first time You have + received notice of non-compliance with this License from such + Contributor, and You become compliant prior to 30 days after Your receipt + of the notice. + +5.2. If You initiate litigation against any entity by asserting a patent + infringement claim (excluding declaratory judgment actions, + counter-claims, and cross-claims) alleging that a Contributor Version + directly or indirectly infringes any patent, then the rights granted to + You by any and all Contributors for the Covered Software under Section + 2.1 of this License shall terminate. + +5.3. In the event of termination under Sections 5.1 or 5.2 above, all end user + license agreements (excluding distributors and resellers) which have been + validly granted by You or Your distributors under this License prior to + termination shall survive termination. + +6. Disclaimer of Warranty + + Covered Software is provided under this License on an "as is" basis, + without warranty of any kind, either expressed, implied, or statutory, + including, without limitation, warranties that the Covered Software is free + of defects, merchantable, fit for a particular purpose or non-infringing. + The entire risk as to the quality and performance of the Covered Software + is with You. Should any Covered Software prove defective in any respect, + You (not any Contributor) assume the cost of any necessary servicing, + repair, or correction. This disclaimer of warranty constitutes an essential + part of this License. No use of any Covered Software is authorized under + this License except under this disclaimer. + +7. Limitation of Liability + + Under no circumstances and under no legal theory, whether tort (including + negligence), contract, or otherwise, shall any Contributor, or anyone who + distributes Covered Software as permitted above, be liable to You for any + direct, indirect, special, incidental, or consequential damages of any + character including, without limitation, damages for lost profits, loss of + goodwill, work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses, even if such party shall have been + informed of the possibility of such damages. This limitation of liability + shall not apply to liability for death or personal injury resulting from + such party's negligence to the extent applicable law prohibits such + limitation. Some jurisdictions do not allow the exclusion or limitation of + incidental or consequential damages, so this exclusion and limitation may + not apply to You. + +8. Litigation + + Any litigation relating to this License may be brought only in the courts + of a jurisdiction where the defendant maintains its principal place of + business and such litigation shall be governed by laws of that + jurisdiction, without reference to its conflict-of-law provisions. Nothing + in this Section shall prevent a party's ability to bring cross-claims or + counter-claims. + +9. Miscellaneous + + This License represents the complete agreement concerning the subject + matter hereof. If any provision of this License is held to be + unenforceable, such provision shall be reformed only to the extent + necessary to make it enforceable. Any law or regulation which provides that + the language of a contract shall be construed against the drafter shall not + be used to construe this License against a Contributor. + + +10. Versions of the License + +10.1. New Versions + + Mozilla Foundation is the license steward. Except as provided in Section + 10.3, no one other than the license steward has the right to modify or + publish new versions of this License. Each version will be given a + distinguishing version number. + +10.2. Effect of New Versions + + You may distribute the Covered Software under the terms of the version + of the License under which You originally received the Covered Software, + or under the terms of any subsequent version published by the license + steward. + +10.3. Modified Versions + + If you create software not governed by this License, and you want to + create a new license for such software, you may create and use a + modified version of this License if you rename the license and remove + any references to the name of the license steward (except to note that + such modified license differs from this License). + +10.4. Distributing Source Code Form that is Incompatible With Secondary + Licenses If You choose to distribute Source Code Form that is + Incompatible With Secondary Licenses under the terms of this version of + the License, the notice described in Exhibit B of this License must be + attached. + +Exhibit A - Source Code Form License Notice + + This Source Code Form is subject to the + terms of the Mozilla Public License, v. + 2.0. If a copy of the MPL was not + distributed with this file, You can + obtain one at + http://mozilla.org/MPL/2.0/. + +If it is not possible or desirable to put the notice in a particular file, +then You may include the notice in a location (such as a LICENSE file in a +relevant directory) where a recipient would be likely to look for such a +notice. + +You may add additional accurate notices of copyright ownership. + +Exhibit B - "Incompatible With Secondary Licenses" Notice + + This Source Code Form is "Incompatible + With Secondary Licenses", as defined by + the Mozilla Public License, v. 2.0. + diff --git a/vendor/github.com/hashicorp/go-immutable-radix/README.md b/vendor/github.com/hashicorp/go-immutable-radix/README.md new file mode 100644 index 00000000..8910fcc0 --- /dev/null +++ b/vendor/github.com/hashicorp/go-immutable-radix/README.md @@ -0,0 +1,41 @@ +go-immutable-radix [![Build Status](https://travis-ci.org/hashicorp/go-immutable-radix.png)](https://travis-ci.org/hashicorp/go-immutable-radix) +========= + +Provides the `iradix` package that implements an immutable [radix tree](http://en.wikipedia.org/wiki/Radix_tree). +The package only provides a single `Tree` implementation, optimized for sparse nodes. + +As a radix tree, it provides the following: + * O(k) operations. In many cases, this can be faster than a hash table since + the hash function is an O(k) operation, and hash tables have very poor cache locality. + * Minimum / Maximum value lookups + * Ordered iteration + +A tree supports using a transaction to batch multiple updates (insert, delete) +in a more efficient manner than performing each operation one at a time. + +For a mutable variant, see [go-radix](https://github.com/armon/go-radix). + +Documentation +============= + +The full documentation is available on [Godoc](http://godoc.org/github.com/hashicorp/go-immutable-radix). + +Example +======= + +Below is a simple example of usage + +```go +// Create a tree +r := iradix.New() +r, _, _ = r.Insert([]byte("foo"), 1) +r, _, _ = r.Insert([]byte("bar"), 2) +r, _, _ = r.Insert([]byte("foobar"), 2) + +// Find the longest prefix match +m, _, _ := r.Root().LongestPrefix([]byte("foozip")) +if string(m) != "foo" { + panic("should be foo") +} +``` + diff --git a/vendor/github.com/hashicorp/go-immutable-radix/edges.go b/vendor/github.com/hashicorp/go-immutable-radix/edges.go new file mode 100644 index 00000000..a6367477 --- /dev/null +++ b/vendor/github.com/hashicorp/go-immutable-radix/edges.go @@ -0,0 +1,21 @@ +package iradix + +import "sort" + +type edges []edge + +func (e edges) Len() int { + return len(e) +} + +func (e edges) Less(i, j int) bool { + return e[i].label < e[j].label +} + +func (e edges) Swap(i, j int) { + e[i], e[j] = e[j], e[i] +} + +func (e edges) Sort() { + sort.Sort(e) +} diff --git a/vendor/github.com/hashicorp/go-immutable-radix/iradix.go b/vendor/github.com/hashicorp/go-immutable-radix/iradix.go new file mode 100644 index 00000000..c7172c40 --- /dev/null +++ b/vendor/github.com/hashicorp/go-immutable-radix/iradix.go @@ -0,0 +1,657 @@ +package iradix + +import ( + "bytes" + "strings" + + "github.com/hashicorp/golang-lru/simplelru" +) + +const ( + // defaultModifiedCache is the default size of the modified node + // cache used per transaction. This is used to cache the updates + // to the nodes near the root, while the leaves do not need to be + // cached. This is important for very large transactions to prevent + // the modified cache from growing to be enormous. This is also used + // to set the max size of the mutation notify maps since those should + // also be bounded in a similar way. + defaultModifiedCache = 8192 +) + +// Tree implements an immutable radix tree. This can be treated as a +// Dictionary abstract data type. The main advantage over a standard +// hash map is prefix-based lookups and ordered iteration. The immutability +// means that it is safe to concurrently read from a Tree without any +// coordination. +type Tree struct { + root *Node + size int +} + +// New returns an empty Tree +func New() *Tree { + t := &Tree{ + root: &Node{ + mutateCh: make(chan struct{}), + }, + } + return t +} + +// Len is used to return the number of elements in the tree +func (t *Tree) Len() int { + return t.size +} + +// Txn is a transaction on the tree. This transaction is applied +// atomically and returns a new tree when committed. A transaction +// is not thread safe, and should only be used by a single goroutine. +type Txn struct { + // root is the modified root for the transaction. + root *Node + + // snap is a snapshot of the root node for use if we have to run the + // slow notify algorithm. + snap *Node + + // size tracks the size of the tree as it is modified during the + // transaction. + size int + + // writable is a cache of writable nodes that have been created during + // the course of the transaction. This allows us to re-use the same + // nodes for further writes and avoid unnecessary copies of nodes that + // have never been exposed outside the transaction. This will only hold + // up to defaultModifiedCache number of entries. + writable *simplelru.LRU + + // trackChannels is used to hold channels that need to be notified to + // signal mutation of the tree. This will only hold up to + // defaultModifiedCache number of entries, after which we will set the + // trackOverflow flag, which will cause us to use a more expensive + // algorithm to perform the notifications. Mutation tracking is only + // performed if trackMutate is true. + trackChannels map[chan struct{}]struct{} + trackOverflow bool + trackMutate bool +} + +// Txn starts a new transaction that can be used to mutate the tree +func (t *Tree) Txn() *Txn { + txn := &Txn{ + root: t.root, + snap: t.root, + size: t.size, + } + return txn +} + +// TrackMutate can be used to toggle if mutations are tracked. If this is enabled +// then notifications will be issued for affected internal nodes and leaves when +// the transaction is committed. +func (t *Txn) TrackMutate(track bool) { + t.trackMutate = track +} + +// trackChannel safely attempts to track the given mutation channel, setting the +// overflow flag if we can no longer track any more. This limits the amount of +// state that will accumulate during a transaction and we have a slower algorithm +// to switch to if we overflow. +func (t *Txn) trackChannel(ch chan struct{}) { + // In overflow, make sure we don't store any more objects. + if t.trackOverflow { + return + } + + // If this would overflow the state we reject it and set the flag (since + // we aren't tracking everything that's required any longer). + if len(t.trackChannels) >= defaultModifiedCache { + // Mark that we are in the overflow state + t.trackOverflow = true + + // Clear the map so that the channels can be garbage collected. It is + // safe to do this since we have already overflowed and will be using + // the slow notify algorithm. + t.trackChannels = nil + return + } + + // Create the map on the fly when we need it. + if t.trackChannels == nil { + t.trackChannels = make(map[chan struct{}]struct{}) + } + + // Otherwise we are good to track it. + t.trackChannels[ch] = struct{}{} +} + +// writeNode returns a node to be modified, if the current node has already been +// modified during the course of the transaction, it is used in-place. Set +// forLeafUpdate to true if you are getting a write node to update the leaf, +// which will set leaf mutation tracking appropriately as well. +func (t *Txn) writeNode(n *Node, forLeafUpdate bool) *Node { + // Ensure the writable set exists. + if t.writable == nil { + lru, err := simplelru.NewLRU(defaultModifiedCache, nil) + if err != nil { + panic(err) + } + t.writable = lru + } + + // If this node has already been modified, we can continue to use it + // during this transaction. We know that we don't need to track it for + // a node update since the node is writable, but if this is for a leaf + // update we track it, in case the initial write to this node didn't + // update the leaf. + if _, ok := t.writable.Get(n); ok { + if t.trackMutate && forLeafUpdate && n.leaf != nil { + t.trackChannel(n.leaf.mutateCh) + } + return n + } + + // Mark this node as being mutated. + if t.trackMutate { + t.trackChannel(n.mutateCh) + } + + // Mark its leaf as being mutated, if appropriate. + if t.trackMutate && forLeafUpdate && n.leaf != nil { + t.trackChannel(n.leaf.mutateCh) + } + + // Copy the existing node. If you have set forLeafUpdate it will be + // safe to replace this leaf with another after you get your node for + // writing. You MUST replace it, because the channel associated with + // this leaf will be closed when this transaction is committed. + nc := &Node{ + mutateCh: make(chan struct{}), + leaf: n.leaf, + } + if n.prefix != nil { + nc.prefix = make([]byte, len(n.prefix)) + copy(nc.prefix, n.prefix) + } + if len(n.edges) != 0 { + nc.edges = make([]edge, len(n.edges)) + copy(nc.edges, n.edges) + } + + // Mark this node as writable. + t.writable.Add(nc, nil) + return nc +} + +// Visit all the nodes in the tree under n, and add their mutateChannels to the transaction +// Returns the size of the subtree visited +func (t *Txn) trackChannelsAndCount(n *Node) int { + // Count only leaf nodes + leaves := 0 + if n.leaf != nil { + leaves = 1 + } + // Mark this node as being mutated. + if t.trackMutate { + t.trackChannel(n.mutateCh) + } + + // Mark its leaf as being mutated, if appropriate. + if t.trackMutate && n.leaf != nil { + t.trackChannel(n.leaf.mutateCh) + } + + // Recurse on the children + for _, e := range n.edges { + leaves += t.trackChannelsAndCount(e.node) + } + return leaves +} + +// mergeChild is called to collapse the given node with its child. This is only +// called when the given node is not a leaf and has a single edge. +func (t *Txn) mergeChild(n *Node) { + // Mark the child node as being mutated since we are about to abandon + // it. We don't need to mark the leaf since we are retaining it if it + // is there. + e := n.edges[0] + child := e.node + if t.trackMutate { + t.trackChannel(child.mutateCh) + } + + // Merge the nodes. + n.prefix = concat(n.prefix, child.prefix) + n.leaf = child.leaf + if len(child.edges) != 0 { + n.edges = make([]edge, len(child.edges)) + copy(n.edges, child.edges) + } else { + n.edges = nil + } +} + +// insert does a recursive insertion +func (t *Txn) insert(n *Node, k, search []byte, v interface{}) (*Node, interface{}, bool) { + // Handle key exhaustion + if len(search) == 0 { + var oldVal interface{} + didUpdate := false + if n.isLeaf() { + oldVal = n.leaf.val + didUpdate = true + } + + nc := t.writeNode(n, true) + nc.leaf = &leafNode{ + mutateCh: make(chan struct{}), + key: k, + val: v, + } + return nc, oldVal, didUpdate + } + + // Look for the edge + idx, child := n.getEdge(search[0]) + + // No edge, create one + if child == nil { + e := edge{ + label: search[0], + node: &Node{ + mutateCh: make(chan struct{}), + leaf: &leafNode{ + mutateCh: make(chan struct{}), + key: k, + val: v, + }, + prefix: search, + }, + } + nc := t.writeNode(n, false) + nc.addEdge(e) + return nc, nil, false + } + + // Determine longest prefix of the search key on match + commonPrefix := longestPrefix(search, child.prefix) + if commonPrefix == len(child.prefix) { + search = search[commonPrefix:] + newChild, oldVal, didUpdate := t.insert(child, k, search, v) + if newChild != nil { + nc := t.writeNode(n, false) + nc.edges[idx].node = newChild + return nc, oldVal, didUpdate + } + return nil, oldVal, didUpdate + } + + // Split the node + nc := t.writeNode(n, false) + splitNode := &Node{ + mutateCh: make(chan struct{}), + prefix: search[:commonPrefix], + } + nc.replaceEdge(edge{ + label: search[0], + node: splitNode, + }) + + // Restore the existing child node + modChild := t.writeNode(child, false) + splitNode.addEdge(edge{ + label: modChild.prefix[commonPrefix], + node: modChild, + }) + modChild.prefix = modChild.prefix[commonPrefix:] + + // Create a new leaf node + leaf := &leafNode{ + mutateCh: make(chan struct{}), + key: k, + val: v, + } + + // If the new key is a subset, add to to this node + search = search[commonPrefix:] + if len(search) == 0 { + splitNode.leaf = leaf + return nc, nil, false + } + + // Create a new edge for the node + splitNode.addEdge(edge{ + label: search[0], + node: &Node{ + mutateCh: make(chan struct{}), + leaf: leaf, + prefix: search, + }, + }) + return nc, nil, false +} + +// delete does a recursive deletion +func (t *Txn) delete(parent, n *Node, search []byte) (*Node, *leafNode) { + // Check for key exhaustion + if len(search) == 0 { + if !n.isLeaf() { + return nil, nil + } + + // Remove the leaf node + nc := t.writeNode(n, true) + nc.leaf = nil + + // Check if this node should be merged + if n != t.root && len(nc.edges) == 1 { + t.mergeChild(nc) + } + return nc, n.leaf + } + + // Look for an edge + label := search[0] + idx, child := n.getEdge(label) + if child == nil || !bytes.HasPrefix(search, child.prefix) { + return nil, nil + } + + // Consume the search prefix + search = search[len(child.prefix):] + newChild, leaf := t.delete(n, child, search) + if newChild == nil { + return nil, nil + } + + // Copy this node. WATCH OUT - it's safe to pass "false" here because we + // will only ADD a leaf via nc.mergeChild() if there isn't one due to + // the !nc.isLeaf() check in the logic just below. This is pretty subtle, + // so be careful if you change any of the logic here. + nc := t.writeNode(n, false) + + // Delete the edge if the node has no edges + if newChild.leaf == nil && len(newChild.edges) == 0 { + nc.delEdge(label) + if n != t.root && len(nc.edges) == 1 && !nc.isLeaf() { + t.mergeChild(nc) + } + } else { + nc.edges[idx].node = newChild + } + return nc, leaf +} + +// delete does a recursive deletion +func (t *Txn) deletePrefix(parent, n *Node, search []byte) (*Node, int) { + // Check for key exhaustion + if len(search) == 0 { + nc := t.writeNode(n, true) + if n.isLeaf() { + nc.leaf = nil + } + nc.edges = nil + return nc, t.trackChannelsAndCount(n) + } + + // Look for an edge + label := search[0] + idx, child := n.getEdge(label) + // We make sure that either the child node's prefix starts with the search term, or the search term starts with the child node's prefix + // Need to do both so that we can delete prefixes that don't correspond to any node in the tree + if child == nil || (!bytes.HasPrefix(child.prefix, search) && !bytes.HasPrefix(search, child.prefix)) { + return nil, 0 + } + + // Consume the search prefix + if len(child.prefix) > len(search) { + search = []byte("") + } else { + search = search[len(child.prefix):] + } + newChild, numDeletions := t.deletePrefix(n, child, search) + if newChild == nil { + return nil, 0 + } + // Copy this node. WATCH OUT - it's safe to pass "false" here because we + // will only ADD a leaf via nc.mergeChild() if there isn't one due to + // the !nc.isLeaf() check in the logic just below. This is pretty subtle, + // so be careful if you change any of the logic here. + + nc := t.writeNode(n, false) + + // Delete the edge if the node has no edges + if newChild.leaf == nil && len(newChild.edges) == 0 { + nc.delEdge(label) + if n != t.root && len(nc.edges) == 1 && !nc.isLeaf() { + t.mergeChild(nc) + } + } else { + nc.edges[idx].node = newChild + } + return nc, numDeletions +} + +// Insert is used to add or update a given key. The return provides +// the previous value and a bool indicating if any was set. +func (t *Txn) Insert(k []byte, v interface{}) (interface{}, bool) { + newRoot, oldVal, didUpdate := t.insert(t.root, k, k, v) + if newRoot != nil { + t.root = newRoot + } + if !didUpdate { + t.size++ + } + return oldVal, didUpdate +} + +// Delete is used to delete a given key. Returns the old value if any, +// and a bool indicating if the key was set. +func (t *Txn) Delete(k []byte) (interface{}, bool) { + newRoot, leaf := t.delete(nil, t.root, k) + if newRoot != nil { + t.root = newRoot + } + if leaf != nil { + t.size-- + return leaf.val, true + } + return nil, false +} + +// DeletePrefix is used to delete an entire subtree that matches the prefix +// This will delete all nodes under that prefix +func (t *Txn) DeletePrefix(prefix []byte) bool { + newRoot, numDeletions := t.deletePrefix(nil, t.root, prefix) + if newRoot != nil { + t.root = newRoot + t.size = t.size - numDeletions + return true + } + return false + +} + +// Root returns the current root of the radix tree within this +// transaction. The root is not safe across insert and delete operations, +// but can be used to read the current state during a transaction. +func (t *Txn) Root() *Node { + return t.root +} + +// Get is used to lookup a specific key, returning +// the value and if it was found +func (t *Txn) Get(k []byte) (interface{}, bool) { + return t.root.Get(k) +} + +// GetWatch is used to lookup a specific key, returning +// the watch channel, value and if it was found +func (t *Txn) GetWatch(k []byte) (<-chan struct{}, interface{}, bool) { + return t.root.GetWatch(k) +} + +// Commit is used to finalize the transaction and return a new tree. If mutation +// tracking is turned on then notifications will also be issued. +func (t *Txn) Commit() *Tree { + nt := t.CommitOnly() + if t.trackMutate { + t.Notify() + } + return nt +} + +// CommitOnly is used to finalize the transaction and return a new tree, but +// does not issue any notifications until Notify is called. +func (t *Txn) CommitOnly() *Tree { + nt := &Tree{t.root, t.size} + t.writable = nil + return nt +} + +// slowNotify does a complete comparison of the before and after trees in order +// to trigger notifications. This doesn't require any additional state but it +// is very expensive to compute. +func (t *Txn) slowNotify() { + snapIter := t.snap.rawIterator() + rootIter := t.root.rawIterator() + for snapIter.Front() != nil || rootIter.Front() != nil { + // If we've exhausted the nodes in the old snapshot, we know + // there's nothing remaining to notify. + if snapIter.Front() == nil { + return + } + snapElem := snapIter.Front() + + // If we've exhausted the nodes in the new root, we know we need + // to invalidate everything that remains in the old snapshot. We + // know from the loop condition there's something in the old + // snapshot. + if rootIter.Front() == nil { + close(snapElem.mutateCh) + if snapElem.isLeaf() { + close(snapElem.leaf.mutateCh) + } + snapIter.Next() + continue + } + + // Do one string compare so we can check the various conditions + // below without repeating the compare. + cmp := strings.Compare(snapIter.Path(), rootIter.Path()) + + // If the snapshot is behind the root, then we must have deleted + // this node during the transaction. + if cmp < 0 { + close(snapElem.mutateCh) + if snapElem.isLeaf() { + close(snapElem.leaf.mutateCh) + } + snapIter.Next() + continue + } + + // If the snapshot is ahead of the root, then we must have added + // this node during the transaction. + if cmp > 0 { + rootIter.Next() + continue + } + + // If we have the same path, then we need to see if we mutated a + // node and possibly the leaf. + rootElem := rootIter.Front() + if snapElem != rootElem { + close(snapElem.mutateCh) + if snapElem.leaf != nil && (snapElem.leaf != rootElem.leaf) { + close(snapElem.leaf.mutateCh) + } + } + snapIter.Next() + rootIter.Next() + } +} + +// Notify is used along with TrackMutate to trigger notifications. This must +// only be done once a transaction is committed via CommitOnly, and it is called +// automatically by Commit. +func (t *Txn) Notify() { + if !t.trackMutate { + return + } + + // If we've overflowed the tracking state we can't use it in any way and + // need to do a full tree compare. + if t.trackOverflow { + t.slowNotify() + } else { + for ch := range t.trackChannels { + close(ch) + } + } + + // Clean up the tracking state so that a re-notify is safe (will trigger + // the else clause above which will be a no-op). + t.trackChannels = nil + t.trackOverflow = false +} + +// Insert is used to add or update a given key. The return provides +// the new tree, previous value and a bool indicating if any was set. +func (t *Tree) Insert(k []byte, v interface{}) (*Tree, interface{}, bool) { + txn := t.Txn() + old, ok := txn.Insert(k, v) + return txn.Commit(), old, ok +} + +// Delete is used to delete a given key. Returns the new tree, +// old value if any, and a bool indicating if the key was set. +func (t *Tree) Delete(k []byte) (*Tree, interface{}, bool) { + txn := t.Txn() + old, ok := txn.Delete(k) + return txn.Commit(), old, ok +} + +// DeletePrefix is used to delete all nodes starting with a given prefix. Returns the new tree, +// and a bool indicating if the prefix matched any nodes +func (t *Tree) DeletePrefix(k []byte) (*Tree, bool) { + txn := t.Txn() + ok := txn.DeletePrefix(k) + return txn.Commit(), ok +} + +// Root returns the root node of the tree which can be used for richer +// query operations. +func (t *Tree) Root() *Node { + return t.root +} + +// Get is used to lookup a specific key, returning +// the value and if it was found +func (t *Tree) Get(k []byte) (interface{}, bool) { + return t.root.Get(k) +} + +// longestPrefix finds the length of the shared prefix +// of two strings +func longestPrefix(k1, k2 []byte) int { + max := len(k1) + if l := len(k2); l < max { + max = l + } + var i int + for i = 0; i < max; i++ { + if k1[i] != k2[i] { + break + } + } + return i +} + +// concat two byte slices, returning a third new copy +func concat(a, b []byte) []byte { + c := make([]byte, len(a)+len(b)) + copy(c, a) + copy(c[len(a):], b) + return c +} diff --git a/vendor/github.com/hashicorp/go-immutable-radix/iradix_test.go b/vendor/github.com/hashicorp/go-immutable-radix/iradix_test.go new file mode 100644 index 00000000..bc9c77c2 --- /dev/null +++ b/vendor/github.com/hashicorp/go-immutable-radix/iradix_test.go @@ -0,0 +1,1496 @@ +package iradix + +import ( + "fmt" + "reflect" + "sort" + "testing" + + "github.com/hashicorp/go-uuid" +) + +func CopyTree(t *Tree) *Tree { + nt := &Tree{ + root: CopyNode(t.root), + size: t.size, + } + return nt +} + +func CopyNode(n *Node) *Node { + nn := &Node{} + if n.mutateCh != nil { + nn.mutateCh = n.mutateCh + } + if n.prefix != nil { + nn.prefix = make([]byte, len(n.prefix)) + copy(nn.prefix, n.prefix) + } + if n.leaf != nil { + nn.leaf = CopyLeaf(n.leaf) + } + if len(n.edges) != 0 { + nn.edges = make([]edge, len(n.edges)) + for idx, edge := range n.edges { + nn.edges[idx].label = edge.label + nn.edges[idx].node = CopyNode(edge.node) + } + } + return nn +} + +func CopyLeaf(l *leafNode) *leafNode { + ll := &leafNode{ + mutateCh: l.mutateCh, + key: l.key, + val: l.val, + } + return ll +} + +func TestRadix_HugeTxn(t *testing.T) { + r := New() + + // Insert way more nodes than the cache can fit + txn1 := r.Txn() + var expect []string + for i := 0; i < defaultModifiedCache*100; i++ { + gen, err := uuid.GenerateUUID() + if err != nil { + t.Fatalf("err: %v", err) + } + txn1.Insert([]byte(gen), i) + expect = append(expect, gen) + } + r = txn1.Commit() + sort.Strings(expect) + + // Collect the output, should be sorted + var out []string + fn := func(k []byte, v interface{}) bool { + out = append(out, string(k)) + return false + } + r.Root().Walk(fn) + + // Verify the match + if len(out) != len(expect) { + t.Fatalf("length mis-match: %d vs %d", len(out), len(expect)) + } + for i := 0; i < len(out); i++ { + if out[i] != expect[i] { + t.Fatalf("mis-match: %v %v", out[i], expect[i]) + } + } +} + +func TestRadix(t *testing.T) { + var min, max string + inp := make(map[string]interface{}) + for i := 0; i < 1000; i++ { + gen, err := uuid.GenerateUUID() + if err != nil { + t.Fatalf("err: %v", err) + } + inp[gen] = i + if gen < min || i == 0 { + min = gen + } + if gen > max || i == 0 { + max = gen + } + } + + r := New() + rCopy := CopyTree(r) + for k, v := range inp { + newR, _, _ := r.Insert([]byte(k), v) + if !reflect.DeepEqual(r, rCopy) { + t.Errorf("r: %#v rc: %#v", r, rCopy) + t.Errorf("r: %#v rc: %#v", r.root, rCopy.root) + t.Fatalf("structure modified %d", newR.Len()) + } + r = newR + rCopy = CopyTree(r) + } + + if r.Len() != len(inp) { + t.Fatalf("bad length: %v %v", r.Len(), len(inp)) + } + + for k, v := range inp { + out, ok := r.Get([]byte(k)) + if !ok { + t.Fatalf("missing key: %v", k) + } + if out != v { + t.Fatalf("value mis-match: %v %v", out, v) + } + } + + // Check min and max + outMin, _, _ := r.Root().Minimum() + if string(outMin) != min { + t.Fatalf("bad minimum: %v %v", outMin, min) + } + outMax, _, _ := r.Root().Maximum() + if string(outMax) != max { + t.Fatalf("bad maximum: %v %v", outMax, max) + } + + // Copy the full tree before delete + orig := r + origCopy := CopyTree(r) + + for k, v := range inp { + tree, out, ok := r.Delete([]byte(k)) + r = tree + if !ok { + t.Fatalf("missing key: %v", k) + } + if out != v { + t.Fatalf("value mis-match: %v %v", out, v) + } + } + if r.Len() != 0 { + t.Fatalf("bad length: %v", r.Len()) + } + + if !reflect.DeepEqual(orig, origCopy) { + t.Fatalf("structure modified") + } +} + +func TestRoot(t *testing.T) { + r := New() + r, _, ok := r.Delete(nil) + if ok { + t.Fatalf("bad") + } + r, _, ok = r.Insert(nil, true) + if ok { + t.Fatalf("bad") + } + val, ok := r.Get(nil) + if !ok || val != true { + t.Fatalf("bad: %v %#v", val) + } + r, val, ok = r.Delete(nil) + if !ok || val != true { + t.Fatalf("bad: %v", val) + } +} + +func TestInsert_UpdateFeedback(t *testing.T) { + r := New() + txn1 := r.Txn() + + for i := 0; i < 10; i++ { + var old interface{} + var didUpdate bool + old, didUpdate = txn1.Insert([]byte("helloworld"), i) + if i == 0 { + if old != nil || didUpdate { + t.Fatalf("bad: %d %v %v", i, old, didUpdate) + } + } else { + if old == nil || old.(int) != i-1 || !didUpdate { + t.Fatalf("bad: %d %v %v", i, old, didUpdate) + } + } + } +} + +func TestDelete(t *testing.T) { + r := New() + s := []string{"", "A", "AB"} + + for _, ss := range s { + r, _, _ = r.Insert([]byte(ss), true) + } + var ok bool + for _, ss := range s { + r, _, ok = r.Delete([]byte(ss)) + if !ok { + t.Fatalf("bad %q", ss) + } + } +} + +func TestDeletePrefix(t *testing.T) { + + type exp struct { + desc string + treeNodes []string + prefix string + expectedOut []string + } + + //various test cases where DeletePrefix should succeed + cases := []exp{ + { + "prefix not a node in tree", + []string{ + "", + "test/test1", + "test/test2", + "test/test3", + "R", + "RA"}, + "test", + []string{ + "", + "R", + "RA", + }, + }, + { + "prefix matches a node in tree", + []string{ + "", + "test", + "test/test1", + "test/test2", + "test/test3", + "test/testAAA", + "R", + "RA", + }, + "test", + []string{ + "", + "R", + "RA", + }, + }, + { + "longer prefix, but prefix is not a node in tree", + []string{ + "", + "test/test1", + "test/test2", + "test/test3", + "test/testAAA", + "R", + "RA", + }, + "test/test", + []string{ + "", + "R", + "RA", + }, + }, + { + "prefix only matches one node", + []string{ + "", + "AB", + "ABC", + "AR", + "R", + "RA", + }, + "AR", + []string{ + "", + "AB", + "ABC", + "R", + "RA", + }, + }, + } + + for _, testCase := range cases { + t.Run(testCase.desc, func(t *testing.T) { + r := New() + for _, ss := range testCase.treeNodes { + r, _, _ = r.Insert([]byte(ss), true) + } + if got, want := r.Len(), len(testCase.treeNodes); got != want { + t.Fatalf("Unexpected tree length after insert, got %d want %d ", got, want) + } + r, ok := r.DeletePrefix([]byte(testCase.prefix)) + if !ok { + t.Fatalf("DeletePrefix should have returned true for tree %v, deleting prefix %v", testCase.treeNodes, testCase.prefix) + } + if got, want := r.Len(), len(testCase.expectedOut); got != want { + t.Fatalf("Bad tree length, got %d want %d tree %v, deleting prefix %v ", got, want, testCase.treeNodes, testCase.prefix) + } + + verifyTree(t, testCase.expectedOut, r) + //Delete a non-existant node + r, ok = r.DeletePrefix([]byte("CCCCC")) + if ok { + t.Fatalf("Expected DeletePrefix to return false ") + } + verifyTree(t, testCase.expectedOut, r) + }) + } +} + +func TestTrackMutate_DeletePrefix(t *testing.T) { + + r := New() + + keys := []string{ + "foo", + "foo/bar/baz", + "foo/baz/bar", + "foo/zip/zap", + "bazbaz", + "zipzap", + } + for _, k := range keys { + r, _, _ = r.Insert([]byte(k), nil) + } + if r.Len() != len(keys) { + t.Fatalf("bad len: %v %v", r.Len(), len(keys)) + } + + rootWatch, _, _ := r.Root().GetWatch(nil) + if rootWatch == nil { + t.Fatalf("Should have returned a watch") + } + + nodeWatch1, _, _ := r.Root().GetWatch([]byte("foo/bar/baz")) + if nodeWatch1 == nil { + t.Fatalf("Should have returned a watch") + } + + nodeWatch2, _, _ := r.Root().GetWatch([]byte("foo/baz/bar")) + if nodeWatch2 == nil { + t.Fatalf("Should have returned a watch") + } + + nodeWatch3, _, _ := r.Root().GetWatch([]byte("foo/zip/zap")) + if nodeWatch3 == nil { + t.Fatalf("Should have returned a watch") + } + + unknownNodeWatch, _, _ := r.Root().GetWatch([]byte("bazbaz")) + if unknownNodeWatch == nil { + t.Fatalf("Should have returned a watch") + } + + // Verify that deleting prefixes triggers the right set of watches + txn := r.Txn() + txn.TrackMutate(true) + ok := txn.DeletePrefix([]byte("foo")) + if !ok { + t.Fatalf("Expected delete prefix to return true") + } + if hasAnyClosedMutateCh(r) { + t.Fatalf("Transaction was not committed, no channel should have been closed") + } + + txn.Commit() + + // Verify that all the leaf nodes we set up watches for above get triggered from the delete prefix call + select { + case <-rootWatch: + default: + t.Fatalf("root watch was not triggered") + } + select { + case <-nodeWatch1: + default: + t.Fatalf("node watch was not triggered") + } + select { + case <-nodeWatch2: + default: + t.Fatalf("node watch was not triggered") + } + select { + case <-nodeWatch3: + default: + t.Fatalf("node watch was not triggered") + } + select { + case <-unknownNodeWatch: + t.Fatalf("Unrelated node watch was triggered during a prefix delete") + default: + } + +} + +func verifyTree(t *testing.T, expected []string, r *Tree) { + root := r.Root() + out := []string{} + fn := func(k []byte, v interface{}) bool { + out = append(out, string(k)) + return false + } + root.Walk(fn) + + if !reflect.DeepEqual(expected, out) { + t.Fatalf("Unexpected contents of tree after delete prefix: expected %v, but got %v", expected, out) + } +} + +func TestLongestPrefix(t *testing.T) { + r := New() + + keys := []string{ + "", + "foo", + "foobar", + "foobarbaz", + "foobarbazzip", + "foozip", + } + for _, k := range keys { + r, _, _ = r.Insert([]byte(k), nil) + } + if r.Len() != len(keys) { + t.Fatalf("bad len: %v %v", r.Len(), len(keys)) + } + + type exp struct { + inp string + out string + } + cases := []exp{ + {"a", ""}, + {"abc", ""}, + {"fo", ""}, + {"foo", "foo"}, + {"foob", "foo"}, + {"foobar", "foobar"}, + {"foobarba", "foobar"}, + {"foobarbaz", "foobarbaz"}, + {"foobarbazzi", "foobarbaz"}, + {"foobarbazzip", "foobarbazzip"}, + {"foozi", "foo"}, + {"foozip", "foozip"}, + {"foozipzap", "foozip"}, + } + root := r.Root() + for _, test := range cases { + m, _, ok := root.LongestPrefix([]byte(test.inp)) + if !ok { + t.Fatalf("no match: %v", test) + } + if string(m) != test.out { + t.Fatalf("mis-match: %v %v", m, test) + } + } +} + +func TestWalkPrefix(t *testing.T) { + r := New() + + keys := []string{ + "foobar", + "foo/bar/baz", + "foo/baz/bar", + "foo/zip/zap", + "zipzap", + } + for _, k := range keys { + r, _, _ = r.Insert([]byte(k), nil) + } + if r.Len() != len(keys) { + t.Fatalf("bad len: %v %v", r.Len(), len(keys)) + } + + type exp struct { + inp string + out []string + } + cases := []exp{ + exp{ + "f", + []string{"foobar", "foo/bar/baz", "foo/baz/bar", "foo/zip/zap"}, + }, + exp{ + "foo", + []string{"foobar", "foo/bar/baz", "foo/baz/bar", "foo/zip/zap"}, + }, + exp{ + "foob", + []string{"foobar"}, + }, + exp{ + "foo/", + []string{"foo/bar/baz", "foo/baz/bar", "foo/zip/zap"}, + }, + exp{ + "foo/b", + []string{"foo/bar/baz", "foo/baz/bar"}, + }, + exp{ + "foo/ba", + []string{"foo/bar/baz", "foo/baz/bar"}, + }, + exp{ + "foo/bar", + []string{"foo/bar/baz"}, + }, + exp{ + "foo/bar/baz", + []string{"foo/bar/baz"}, + }, + exp{ + "foo/bar/bazoo", + []string{}, + }, + exp{ + "z", + []string{"zipzap"}, + }, + } + + root := r.Root() + for _, test := range cases { + out := []string{} + fn := func(k []byte, v interface{}) bool { + out = append(out, string(k)) + return false + } + root.WalkPrefix([]byte(test.inp), fn) + sort.Strings(out) + sort.Strings(test.out) + if !reflect.DeepEqual(out, test.out) { + t.Fatalf("mis-match: %v %v", out, test.out) + } + } +} + +func TestWalkPath(t *testing.T) { + r := New() + + keys := []string{ + "foo", + "foo/bar", + "foo/bar/baz", + "foo/baz/bar", + "foo/zip/zap", + "zipzap", + } + for _, k := range keys { + r, _, _ = r.Insert([]byte(k), nil) + } + if r.Len() != len(keys) { + t.Fatalf("bad len: %v %v", r.Len(), len(keys)) + } + + type exp struct { + inp string + out []string + } + cases := []exp{ + exp{ + "f", + []string{}, + }, + exp{ + "foo", + []string{"foo"}, + }, + exp{ + "foo/", + []string{"foo"}, + }, + exp{ + "foo/ba", + []string{"foo"}, + }, + exp{ + "foo/bar", + []string{"foo", "foo/bar"}, + }, + exp{ + "foo/bar/baz", + []string{"foo", "foo/bar", "foo/bar/baz"}, + }, + exp{ + "foo/bar/bazoo", + []string{"foo", "foo/bar", "foo/bar/baz"}, + }, + exp{ + "z", + []string{}, + }, + } + + root := r.Root() + for _, test := range cases { + out := []string{} + fn := func(k []byte, v interface{}) bool { + out = append(out, string(k)) + return false + } + root.WalkPath([]byte(test.inp), fn) + sort.Strings(out) + sort.Strings(test.out) + if !reflect.DeepEqual(out, test.out) { + t.Fatalf("mis-match: %v %v", out, test.out) + } + } +} + +func TestIteratePrefix(t *testing.T) { + r := New() + + keys := []string{ + "foo/bar/baz", + "foo/baz/bar", + "foo/zip/zap", + "foobar", + "zipzap", + } + for _, k := range keys { + r, _, _ = r.Insert([]byte(k), nil) + } + if r.Len() != len(keys) { + t.Fatalf("bad len: %v %v", r.Len(), len(keys)) + } + + type exp struct { + inp string + out []string + } + cases := []exp{ + exp{ + "", + keys, + }, + exp{ + "f", + []string{ + "foo/bar/baz", + "foo/baz/bar", + "foo/zip/zap", + "foobar", + }, + }, + exp{ + "foo", + []string{ + "foo/bar/baz", + "foo/baz/bar", + "foo/zip/zap", + "foobar", + }, + }, + exp{ + "foob", + []string{"foobar"}, + }, + exp{ + "foo/", + []string{"foo/bar/baz", "foo/baz/bar", "foo/zip/zap"}, + }, + exp{ + "foo/b", + []string{"foo/bar/baz", "foo/baz/bar"}, + }, + exp{ + "foo/ba", + []string{"foo/bar/baz", "foo/baz/bar"}, + }, + exp{ + "foo/bar", + []string{"foo/bar/baz"}, + }, + exp{ + "foo/bar/baz", + []string{"foo/bar/baz"}, + }, + exp{ + "foo/bar/bazoo", + []string{}, + }, + exp{ + "z", + []string{"zipzap"}, + }, + } + + root := r.Root() + for idx, test := range cases { + iter := root.Iterator() + if test.inp != "" { + iter.SeekPrefix([]byte(test.inp)) + } + + // Consume all the keys + out := []string{} + for { + key, _, ok := iter.Next() + if !ok { + break + } + out = append(out, string(key)) + } + if !reflect.DeepEqual(out, test.out) { + t.Fatalf("mis-match: %d %v %v", idx, out, test.out) + } + } +} + +func TestMergeChildNilEdges(t *testing.T) { + r := New() + r, _, _ = r.Insert([]byte("foobar"), 42) + r, _, _ = r.Insert([]byte("foozip"), 43) + r, _, _ = r.Delete([]byte("foobar")) + + root := r.Root() + out := []string{} + fn := func(k []byte, v interface{}) bool { + out = append(out, string(k)) + return false + } + root.Walk(fn) + + expect := []string{"foozip"} + sort.Strings(out) + sort.Strings(expect) + if !reflect.DeepEqual(out, expect) { + t.Fatalf("mis-match: %v %v", out, expect) + } +} + +func TestMergeChildVisibility(t *testing.T) { + r := New() + r, _, _ = r.Insert([]byte("foobar"), 42) + r, _, _ = r.Insert([]byte("foobaz"), 43) + r, _, _ = r.Insert([]byte("foozip"), 10) + + txn1 := r.Txn() + txn2 := r.Txn() + + // Ensure we get the expected value foobar and foobaz + if val, ok := txn1.Get([]byte("foobar")); !ok || val != 42 { + t.Fatalf("bad: %v", val) + } + if val, ok := txn1.Get([]byte("foobaz")); !ok || val != 43 { + t.Fatalf("bad: %v", val) + } + if val, ok := txn2.Get([]byte("foobar")); !ok || val != 42 { + t.Fatalf("bad: %v", val) + } + if val, ok := txn2.Get([]byte("foobaz")); !ok || val != 43 { + t.Fatalf("bad: %v", val) + } + + // Delete of foozip will cause a merge child between the + // "foo" and "ba" nodes. + if val, ok := txn2.Delete([]byte("foozip")); !ok || val != 10 { + t.Fatalf("bad: %v", val) + } + + // Insert of "foobaz" will update the slice of the "fooba" node + // in-place to point to the new "foobaz" node. This in-place update + // will cause the visibility of the update to leak into txn1 (prior + // to the fix). + if val, ok := txn2.Insert([]byte("foobaz"), 44); !ok || val != 43 { + t.Fatalf("bad: %v", val) + } + + // Ensure we get the expected value foobar and foobaz + if val, ok := txn1.Get([]byte("foobar")); !ok || val != 42 { + t.Fatalf("bad: %v", val) + } + if val, ok := txn1.Get([]byte("foobaz")); !ok || val != 43 { + t.Fatalf("bad: %v", val) + } + if val, ok := txn2.Get([]byte("foobar")); !ok || val != 42 { + t.Fatalf("bad: %v", val) + } + if val, ok := txn2.Get([]byte("foobaz")); !ok || val != 44 { + t.Fatalf("bad: %v", val) + } + + // Commit txn2 + r = txn2.Commit() + + // Ensure we get the expected value foobar and foobaz + if val, ok := txn1.Get([]byte("foobar")); !ok || val != 42 { + t.Fatalf("bad: %v", val) + } + if val, ok := txn1.Get([]byte("foobaz")); !ok || val != 43 { + t.Fatalf("bad: %v", val) + } + if val, ok := r.Get([]byte("foobar")); !ok || val != 42 { + t.Fatalf("bad: %v", val) + } + if val, ok := r.Get([]byte("foobaz")); !ok || val != 44 { + t.Fatalf("bad: %v", val) + } +} + +// isClosed returns true if the given channel is closed. +func isClosed(ch chan struct{}) bool { + select { + case <-ch: + return true + default: + return false + } +} + +// hasAnyClosedMutateCh scans the given tree and returns true if there are any +// closed mutate channels on any nodes or leaves. +func hasAnyClosedMutateCh(r *Tree) bool { + for iter := r.root.rawIterator(); iter.Front() != nil; iter.Next() { + n := iter.Front() + if isClosed(n.mutateCh) { + return true + } + if n.isLeaf() && isClosed(n.leaf.mutateCh) { + return true + } + } + return false +} + +func TestTrackMutate_SeekPrefixWatch(t *testing.T) { + for i := 0; i < 3; i++ { + r := New() + + keys := []string{ + "foo/bar/baz", + "foo/baz/bar", + "foo/zip/zap", + "foobar", + "zipzap", + } + for _, k := range keys { + r, _, _ = r.Insert([]byte(k), nil) + } + if r.Len() != len(keys) { + t.Fatalf("bad len: %v %v", r.Len(), len(keys)) + } + + iter := r.Root().Iterator() + rootWatch := iter.SeekPrefixWatch([]byte("nope")) + + iter = r.Root().Iterator() + parentWatch := iter.SeekPrefixWatch([]byte("foo")) + + iter = r.Root().Iterator() + leafWatch := iter.SeekPrefixWatch([]byte("foobar")) + + iter = r.Root().Iterator() + missingWatch := iter.SeekPrefixWatch([]byte("foobarbaz")) + + iter = r.Root().Iterator() + otherWatch := iter.SeekPrefixWatch([]byte("foo/b")) + + // Write to a sub-child should trigger the leaf! + txn := r.Txn() + txn.TrackMutate(true) + txn.Insert([]byte("foobarbaz"), nil) + switch i { + case 0: + r = txn.Commit() + case 1: + r = txn.CommitOnly() + txn.Notify() + default: + r = txn.CommitOnly() + txn.slowNotify() + } + if hasAnyClosedMutateCh(r) { + t.Fatalf("bad") + } + + // Verify root and parent triggered, and leaf affected + select { + case <-rootWatch: + default: + t.Fatalf("bad") + } + select { + case <-parentWatch: + default: + t.Fatalf("bad") + } + select { + case <-leafWatch: + default: + t.Fatalf("bad") + } + select { + case <-missingWatch: + default: + t.Fatalf("bad") + } + select { + case <-otherWatch: + t.Fatalf("bad") + default: + } + + iter = r.Root().Iterator() + rootWatch = iter.SeekPrefixWatch([]byte("nope")) + + iter = r.Root().Iterator() + parentWatch = iter.SeekPrefixWatch([]byte("foo")) + + iter = r.Root().Iterator() + leafWatch = iter.SeekPrefixWatch([]byte("foobar")) + + iter = r.Root().Iterator() + missingWatch = iter.SeekPrefixWatch([]byte("foobarbaz")) + + // Delete to a sub-child should trigger the leaf! + txn = r.Txn() + txn.TrackMutate(true) + txn.Delete([]byte("foobarbaz")) + switch i { + case 0: + r = txn.Commit() + case 1: + r = txn.CommitOnly() + txn.Notify() + default: + r = txn.CommitOnly() + txn.slowNotify() + } + if hasAnyClosedMutateCh(r) { + t.Fatalf("bad") + } + + // Verify root and parent triggered, and leaf affected + select { + case <-rootWatch: + default: + t.Fatalf("bad") + } + select { + case <-parentWatch: + default: + t.Fatalf("bad") + } + select { + case <-leafWatch: + default: + t.Fatalf("bad") + } + select { + case <-missingWatch: + default: + t.Fatalf("bad") + } + select { + case <-otherWatch: + t.Fatalf("bad") + default: + } + } +} + +func TestTrackMutate_GetWatch(t *testing.T) { + for i := 0; i < 3; i++ { + r := New() + + keys := []string{ + "foo/bar/baz", + "foo/baz/bar", + "foo/zip/zap", + "foobar", + "zipzap", + } + for _, k := range keys { + r, _, _ = r.Insert([]byte(k), nil) + } + if r.Len() != len(keys) { + t.Fatalf("bad len: %v %v", r.Len(), len(keys)) + } + + rootWatch, _, ok := r.Root().GetWatch(nil) + if rootWatch == nil { + t.Fatalf("bad") + } + + parentWatch, _, ok := r.Root().GetWatch([]byte("foo")) + if parentWatch == nil { + t.Fatalf("bad") + } + + leafWatch, _, ok := r.Root().GetWatch([]byte("foobar")) + if !ok { + t.Fatalf("should be found") + } + if leafWatch == nil { + t.Fatalf("bad") + } + + otherWatch, _, ok := r.Root().GetWatch([]byte("foo/b")) + if otherWatch == nil { + t.Fatalf("bad") + } + + // Write to a sub-child should not trigger the leaf! + txn := r.Txn() + txn.TrackMutate(true) + txn.Insert([]byte("foobarbaz"), nil) + switch i { + case 0: + r = txn.Commit() + case 1: + r = txn.CommitOnly() + txn.Notify() + default: + r = txn.CommitOnly() + txn.slowNotify() + } + if hasAnyClosedMutateCh(r) { + t.Fatalf("bad") + } + + // Verify root and parent triggered, not leaf affected + select { + case <-rootWatch: + default: + t.Fatalf("bad") + } + select { + case <-parentWatch: + default: + t.Fatalf("bad") + } + select { + case <-leafWatch: + t.Fatalf("bad") + default: + } + select { + case <-otherWatch: + t.Fatalf("bad") + default: + } + + // Setup new watchers + rootWatch, _, ok = r.Root().GetWatch(nil) + if rootWatch == nil { + t.Fatalf("bad") + } + + parentWatch, _, ok = r.Root().GetWatch([]byte("foo")) + if parentWatch == nil { + t.Fatalf("bad") + } + + // Write to a exactly leaf should trigger the leaf! + txn = r.Txn() + txn.TrackMutate(true) + txn.Insert([]byte("foobar"), nil) + switch i { + case 0: + r = txn.Commit() + case 1: + r = txn.CommitOnly() + txn.Notify() + default: + r = txn.CommitOnly() + txn.slowNotify() + } + if hasAnyClosedMutateCh(r) { + t.Fatalf("bad") + } + + select { + case <-rootWatch: + default: + t.Fatalf("bad") + } + select { + case <-parentWatch: + default: + t.Fatalf("bad") + } + select { + case <-leafWatch: + default: + t.Fatalf("bad") + } + select { + case <-otherWatch: + t.Fatalf("bad") + default: + } + + // Setup all the watchers again + rootWatch, _, ok = r.Root().GetWatch(nil) + if rootWatch == nil { + t.Fatalf("bad") + } + + parentWatch, _, ok = r.Root().GetWatch([]byte("foo")) + if parentWatch == nil { + t.Fatalf("bad") + } + + leafWatch, _, ok = r.Root().GetWatch([]byte("foobar")) + if !ok { + t.Fatalf("should be found") + } + if leafWatch == nil { + t.Fatalf("bad") + } + + // Delete to a sub-child should not trigger the leaf! + txn = r.Txn() + txn.TrackMutate(true) + txn.Delete([]byte("foobarbaz")) + switch i { + case 0: + r = txn.Commit() + case 1: + r = txn.CommitOnly() + txn.Notify() + default: + r = txn.CommitOnly() + txn.slowNotify() + } + if hasAnyClosedMutateCh(r) { + t.Fatalf("bad") + } + + // Verify root and parent triggered, not leaf affected + select { + case <-rootWatch: + default: + t.Fatalf("bad") + } + select { + case <-parentWatch: + default: + t.Fatalf("bad") + } + select { + case <-leafWatch: + t.Fatalf("bad") + default: + } + select { + case <-otherWatch: + t.Fatalf("bad") + default: + } + + // Setup new watchers + rootWatch, _, ok = r.Root().GetWatch(nil) + if rootWatch == nil { + t.Fatalf("bad") + } + + parentWatch, _, ok = r.Root().GetWatch([]byte("foo")) + if parentWatch == nil { + t.Fatalf("bad") + } + + // Write to a exactly leaf should trigger the leaf! + txn = r.Txn() + txn.TrackMutate(true) + txn.Delete([]byte("foobar")) + switch i { + case 0: + r = txn.Commit() + case 1: + r = txn.CommitOnly() + txn.Notify() + default: + r = txn.CommitOnly() + txn.slowNotify() + } + if hasAnyClosedMutateCh(r) { + t.Fatalf("bad") + } + + select { + case <-rootWatch: + default: + t.Fatalf("bad") + } + select { + case <-parentWatch: + default: + t.Fatalf("bad") + } + select { + case <-leafWatch: + default: + t.Fatalf("bad") + } + select { + case <-otherWatch: + t.Fatalf("bad") + default: + } + } +} + +func TestTrackMutate_HugeTxn(t *testing.T) { + r := New() + + keys := []string{ + "foo/bar/baz", + "foo/baz/bar", + "foo/zip/zap", + "foobar", + "nochange", + } + for i := 0; i < defaultModifiedCache; i++ { + key := fmt.Sprintf("aaa%d", i) + r, _, _ = r.Insert([]byte(key), nil) + } + for _, k := range keys { + r, _, _ = r.Insert([]byte(k), nil) + } + for i := 0; i < defaultModifiedCache; i++ { + key := fmt.Sprintf("zzz%d", i) + r, _, _ = r.Insert([]byte(key), nil) + } + if r.Len() != len(keys)+2*defaultModifiedCache { + t.Fatalf("bad len: %v %v", r.Len(), len(keys)) + } + + rootWatch, _, ok := r.Root().GetWatch(nil) + if rootWatch == nil { + t.Fatalf("bad") + } + + parentWatch, _, ok := r.Root().GetWatch([]byte("foo")) + if parentWatch == nil { + t.Fatalf("bad") + } + + leafWatch, _, ok := r.Root().GetWatch([]byte("foobar")) + if !ok { + t.Fatalf("should be found") + } + if leafWatch == nil { + t.Fatalf("bad") + } + + nopeWatch, _, ok := r.Root().GetWatch([]byte("nochange")) + if !ok { + t.Fatalf("should be found") + } + if nopeWatch == nil { + t.Fatalf("bad") + } + + beforeWatch, _, ok := r.Root().GetWatch([]byte("aaa123")) + if beforeWatch == nil { + t.Fatalf("bad") + } + + afterWatch, _, ok := r.Root().GetWatch([]byte("zzz123")) + if afterWatch == nil { + t.Fatalf("bad") + } + + // Start the transaction. + txn := r.Txn() + txn.TrackMutate(true) + + // Add new nodes on both sides of the tree and delete enough nodes to + // overflow the tracking. + txn.Insert([]byte("aaa"), nil) + for i := 0; i < defaultModifiedCache; i++ { + key := fmt.Sprintf("aaa%d", i) + txn.Delete([]byte(key)) + } + for i := 0; i < defaultModifiedCache; i++ { + key := fmt.Sprintf("zzz%d", i) + txn.Delete([]byte(key)) + } + txn.Insert([]byte("zzz"), nil) + + // Hit the leaf, and add a child so we make multiple mutations to the + // same node. + txn.Insert([]byte("foobar"), nil) + txn.Insert([]byte("foobarbaz"), nil) + + // Commit and make sure we overflowed but didn't take on extra stuff. + r = txn.CommitOnly() + if !txn.trackOverflow || txn.trackChannels != nil { + t.Fatalf("bad") + } + + // Now do the trigger. + txn.Notify() + + // Make sure no closed channels escaped the transaction. + if hasAnyClosedMutateCh(r) { + t.Fatalf("bad") + } + + // Verify the watches fired as expected. + select { + case <-rootWatch: + default: + t.Fatalf("bad") + } + select { + case <-parentWatch: + default: + t.Fatalf("bad") + } + select { + case <-leafWatch: + default: + t.Fatalf("bad") + } + select { + case <-nopeWatch: + t.Fatalf("bad") + default: + } + select { + case <-beforeWatch: + default: + t.Fatalf("bad") + } + select { + case <-afterWatch: + default: + t.Fatalf("bad") + } +} + +func TestTrackMutate_mergeChild(t *testing.T) { + // This case does a delete of the "acb" leaf, which causes the "aca" + // leaf to get merged with the old "ac" node: + // + // [root] [root] + // |a |a + // [node] [node] + // b/ \c b/ \c + // (ab) [node] (ab) (aca) + // a/ \b + // (aca) (acb) + // + for i := 0; i < 3; i++ { + r := New() + r, _, _ = r.Insert([]byte("ab"), nil) + r, _, _ = r.Insert([]byte("aca"), nil) + r, _, _ = r.Insert([]byte("acb"), nil) + snapIter := r.root.rawIterator() + + // Run through all notification methods as there were bugs in + // both that affected these operations. The slowNotify path + // would detect copied but otherwise identical leaves as changed + // and wrongly close channels. The normal path would fail to + // notify on a child node that had been merged. + txn := r.Txn() + txn.TrackMutate(true) + txn.Delete([]byte("acb")) + switch i { + case 0: + r = txn.Commit() + case 1: + r = txn.CommitOnly() + txn.Notify() + default: + r = txn.CommitOnly() + txn.slowNotify() + } + if hasAnyClosedMutateCh(r) { + t.Fatalf("bad") + } + + // Run through the old tree and make sure the exact channels we + // expected were closed. + for ; snapIter.Front() != nil; snapIter.Next() { + n := snapIter.Front() + path := snapIter.Path() + switch path { + case "", "a", "ac": // parent nodes all change + if !isClosed(n.mutateCh) || n.leaf != nil { + t.Fatalf("bad") + } + case "ab": // unrelated node / leaf sees no change + if isClosed(n.mutateCh) || isClosed(n.leaf.mutateCh) { + t.Fatalf("bad") + } + case "aca": // this node gets merged, but the leaf doesn't change + if !isClosed(n.mutateCh) || isClosed(n.leaf.mutateCh) { + t.Fatalf("bad") + } + case "acb": // this node / leaf gets deleted + if !isClosed(n.mutateCh) || !isClosed(n.leaf.mutateCh) { + t.Fatalf("bad") + } + default: + t.Fatalf("bad: %s", path) + } + } + } +} + +func TestTrackMutate_cachedNodeChange(t *testing.T) { + // This case does a delete of the "acb" leaf, which causes the "aca" + // leaf to get merged with the old "ac" node: + // + // [root] [root] + // |a |a + // [node] [node] + // b/ \c b/ \c + // (ab) [node] (ab) (aca*) <- this leaf gets modified + // a/ \b post-merge + // (aca) (acb) + // + // Then it makes a modification to the "aca" leaf on a node that will + // be in the cache, so this makes sure that the leaf watch fires. + for i := 0; i < 3; i++ { + r := New() + r, _, _ = r.Insert([]byte("ab"), nil) + r, _, _ = r.Insert([]byte("aca"), nil) + r, _, _ = r.Insert([]byte("acb"), nil) + snapIter := r.root.rawIterator() + + txn := r.Txn() + txn.TrackMutate(true) + txn.Delete([]byte("acb")) + txn.Insert([]byte("aca"), nil) + switch i { + case 0: + r = txn.Commit() + case 1: + r = txn.CommitOnly() + txn.Notify() + default: + r = txn.CommitOnly() + txn.slowNotify() + } + if hasAnyClosedMutateCh(r) { + t.Fatalf("bad") + } + + // Run through the old tree and make sure the exact channels we + // expected were closed. + for ; snapIter.Front() != nil; snapIter.Next() { + n := snapIter.Front() + path := snapIter.Path() + switch path { + case "", "a", "ac": // parent nodes all change + if !isClosed(n.mutateCh) || n.leaf != nil { + t.Fatalf("bad") + } + case "ab": // unrelated node / leaf sees no change + if isClosed(n.mutateCh) || isClosed(n.leaf.mutateCh) { + t.Fatalf("bad") + } + case "aca": // merge changes the node, then we update the leaf + if !isClosed(n.mutateCh) || !isClosed(n.leaf.mutateCh) { + t.Fatalf("bad") + } + case "acb": // this node / leaf gets deleted + if !isClosed(n.mutateCh) || !isClosed(n.leaf.mutateCh) { + t.Fatalf("bad") + } + default: + t.Fatalf("bad: %s", path) + } + } + } +} diff --git a/vendor/github.com/hashicorp/go-immutable-radix/iter.go b/vendor/github.com/hashicorp/go-immutable-radix/iter.go new file mode 100644 index 00000000..9815e025 --- /dev/null +++ b/vendor/github.com/hashicorp/go-immutable-radix/iter.go @@ -0,0 +1,91 @@ +package iradix + +import "bytes" + +// Iterator is used to iterate over a set of nodes +// in pre-order +type Iterator struct { + node *Node + stack []edges +} + +// SeekPrefixWatch is used to seek the iterator to a given prefix +// and returns the watch channel of the finest granularity +func (i *Iterator) SeekPrefixWatch(prefix []byte) (watch <-chan struct{}) { + // Wipe the stack + i.stack = nil + n := i.node + watch = n.mutateCh + search := prefix + for { + // Check for key exhaution + if len(search) == 0 { + i.node = n + return + } + + // Look for an edge + _, n = n.getEdge(search[0]) + if n == nil { + i.node = nil + return + } + + // Update to the finest granularity as the search makes progress + watch = n.mutateCh + + // Consume the search prefix + if bytes.HasPrefix(search, n.prefix) { + search = search[len(n.prefix):] + + } else if bytes.HasPrefix(n.prefix, search) { + i.node = n + return + } else { + i.node = nil + return + } + } +} + +// SeekPrefix is used to seek the iterator to a given prefix +func (i *Iterator) SeekPrefix(prefix []byte) { + i.SeekPrefixWatch(prefix) +} + +// Next returns the next node in order +func (i *Iterator) Next() ([]byte, interface{}, bool) { + // Initialize our stack if needed + if i.stack == nil && i.node != nil { + i.stack = []edges{ + edges{ + edge{node: i.node}, + }, + } + } + + for len(i.stack) > 0 { + // Inspect the last element of the stack + n := len(i.stack) + last := i.stack[n-1] + elem := last[0].node + + // Update the stack + if len(last) > 1 { + i.stack[n-1] = last[1:] + } else { + i.stack = i.stack[:n-1] + } + + // Push the edges onto the frontier + if len(elem.edges) > 0 { + i.stack = append(i.stack, elem.edges) + } + + // Return the leaf values if any + if elem.leaf != nil { + return elem.leaf.key, elem.leaf.val, true + } + } + return nil, nil, false +} diff --git a/vendor/github.com/hashicorp/go-immutable-radix/node.go b/vendor/github.com/hashicorp/go-immutable-radix/node.go new file mode 100644 index 00000000..7a065e7a --- /dev/null +++ b/vendor/github.com/hashicorp/go-immutable-radix/node.go @@ -0,0 +1,292 @@ +package iradix + +import ( + "bytes" + "sort" +) + +// WalkFn is used when walking the tree. Takes a +// key and value, returning if iteration should +// be terminated. +type WalkFn func(k []byte, v interface{}) bool + +// leafNode is used to represent a value +type leafNode struct { + mutateCh chan struct{} + key []byte + val interface{} +} + +// edge is used to represent an edge node +type edge struct { + label byte + node *Node +} + +// Node is an immutable node in the radix tree +type Node struct { + // mutateCh is closed if this node is modified + mutateCh chan struct{} + + // leaf is used to store possible leaf + leaf *leafNode + + // prefix is the common prefix we ignore + prefix []byte + + // Edges should be stored in-order for iteration. + // We avoid a fully materialized slice to save memory, + // since in most cases we expect to be sparse + edges edges +} + +func (n *Node) isLeaf() bool { + return n.leaf != nil +} + +func (n *Node) addEdge(e edge) { + num := len(n.edges) + idx := sort.Search(num, func(i int) bool { + return n.edges[i].label >= e.label + }) + n.edges = append(n.edges, e) + if idx != num { + copy(n.edges[idx+1:], n.edges[idx:num]) + n.edges[idx] = e + } +} + +func (n *Node) replaceEdge(e edge) { + num := len(n.edges) + idx := sort.Search(num, func(i int) bool { + return n.edges[i].label >= e.label + }) + if idx < num && n.edges[idx].label == e.label { + n.edges[idx].node = e.node + return + } + panic("replacing missing edge") +} + +func (n *Node) getEdge(label byte) (int, *Node) { + num := len(n.edges) + idx := sort.Search(num, func(i int) bool { + return n.edges[i].label >= label + }) + if idx < num && n.edges[idx].label == label { + return idx, n.edges[idx].node + } + return -1, nil +} + +func (n *Node) delEdge(label byte) { + num := len(n.edges) + idx := sort.Search(num, func(i int) bool { + return n.edges[i].label >= label + }) + if idx < num && n.edges[idx].label == label { + copy(n.edges[idx:], n.edges[idx+1:]) + n.edges[len(n.edges)-1] = edge{} + n.edges = n.edges[:len(n.edges)-1] + } +} + +func (n *Node) GetWatch(k []byte) (<-chan struct{}, interface{}, bool) { + search := k + watch := n.mutateCh + for { + // Check for key exhaustion + if len(search) == 0 { + if n.isLeaf() { + return n.leaf.mutateCh, n.leaf.val, true + } + break + } + + // Look for an edge + _, n = n.getEdge(search[0]) + if n == nil { + break + } + + // Update to the finest granularity as the search makes progress + watch = n.mutateCh + + // Consume the search prefix + if bytes.HasPrefix(search, n.prefix) { + search = search[len(n.prefix):] + } else { + break + } + } + return watch, nil, false +} + +func (n *Node) Get(k []byte) (interface{}, bool) { + _, val, ok := n.GetWatch(k) + return val, ok +} + +// LongestPrefix is like Get, but instead of an +// exact match, it will return the longest prefix match. +func (n *Node) LongestPrefix(k []byte) ([]byte, interface{}, bool) { + var last *leafNode + search := k + for { + // Look for a leaf node + if n.isLeaf() { + last = n.leaf + } + + // Check for key exhaution + if len(search) == 0 { + break + } + + // Look for an edge + _, n = n.getEdge(search[0]) + if n == nil { + break + } + + // Consume the search prefix + if bytes.HasPrefix(search, n.prefix) { + search = search[len(n.prefix):] + } else { + break + } + } + if last != nil { + return last.key, last.val, true + } + return nil, nil, false +} + +// Minimum is used to return the minimum value in the tree +func (n *Node) Minimum() ([]byte, interface{}, bool) { + for { + if n.isLeaf() { + return n.leaf.key, n.leaf.val, true + } + if len(n.edges) > 0 { + n = n.edges[0].node + } else { + break + } + } + return nil, nil, false +} + +// Maximum is used to return the maximum value in the tree +func (n *Node) Maximum() ([]byte, interface{}, bool) { + for { + if num := len(n.edges); num > 0 { + n = n.edges[num-1].node + continue + } + if n.isLeaf() { + return n.leaf.key, n.leaf.val, true + } else { + break + } + } + return nil, nil, false +} + +// Iterator is used to return an iterator at +// the given node to walk the tree +func (n *Node) Iterator() *Iterator { + return &Iterator{node: n} +} + +// rawIterator is used to return a raw iterator at the given node to walk the +// tree. +func (n *Node) rawIterator() *rawIterator { + iter := &rawIterator{node: n} + iter.Next() + return iter +} + +// Walk is used to walk the tree +func (n *Node) Walk(fn WalkFn) { + recursiveWalk(n, fn) +} + +// WalkPrefix is used to walk the tree under a prefix +func (n *Node) WalkPrefix(prefix []byte, fn WalkFn) { + search := prefix + for { + // Check for key exhaution + if len(search) == 0 { + recursiveWalk(n, fn) + return + } + + // Look for an edge + _, n = n.getEdge(search[0]) + if n == nil { + break + } + + // Consume the search prefix + if bytes.HasPrefix(search, n.prefix) { + search = search[len(n.prefix):] + + } else if bytes.HasPrefix(n.prefix, search) { + // Child may be under our search prefix + recursiveWalk(n, fn) + return + } else { + break + } + } +} + +// WalkPath is used to walk the tree, but only visiting nodes +// from the root down to a given leaf. Where WalkPrefix walks +// all the entries *under* the given prefix, this walks the +// entries *above* the given prefix. +func (n *Node) WalkPath(path []byte, fn WalkFn) { + search := path + for { + // Visit the leaf values if any + if n.leaf != nil && fn(n.leaf.key, n.leaf.val) { + return + } + + // Check for key exhaution + if len(search) == 0 { + return + } + + // Look for an edge + _, n = n.getEdge(search[0]) + if n == nil { + return + } + + // Consume the search prefix + if bytes.HasPrefix(search, n.prefix) { + search = search[len(n.prefix):] + } else { + break + } + } +} + +// recursiveWalk is used to do a pre-order walk of a node +// recursively. Returns true if the walk should be aborted +func recursiveWalk(n *Node, fn WalkFn) bool { + // Visit the leaf values if any + if n.leaf != nil && fn(n.leaf.key, n.leaf.val) { + return true + } + + // Recurse on the children + for _, e := range n.edges { + if recursiveWalk(e.node, fn) { + return true + } + } + return false +} diff --git a/vendor/github.com/hashicorp/go-immutable-radix/raw_iter.go b/vendor/github.com/hashicorp/go-immutable-radix/raw_iter.go new file mode 100644 index 00000000..04814c13 --- /dev/null +++ b/vendor/github.com/hashicorp/go-immutable-radix/raw_iter.go @@ -0,0 +1,78 @@ +package iradix + +// rawIterator visits each of the nodes in the tree, even the ones that are not +// leaves. It keeps track of the effective path (what a leaf at a given node +// would be called), which is useful for comparing trees. +type rawIterator struct { + // node is the starting node in the tree for the iterator. + node *Node + + // stack keeps track of edges in the frontier. + stack []rawStackEntry + + // pos is the current position of the iterator. + pos *Node + + // path is the effective path of the current iterator position, + // regardless of whether the current node is a leaf. + path string +} + +// rawStackEntry is used to keep track of the cumulative common path as well as +// its associated edges in the frontier. +type rawStackEntry struct { + path string + edges edges +} + +// Front returns the current node that has been iterated to. +func (i *rawIterator) Front() *Node { + return i.pos +} + +// Path returns the effective path of the current node, even if it's not actually +// a leaf. +func (i *rawIterator) Path() string { + return i.path +} + +// Next advances the iterator to the next node. +func (i *rawIterator) Next() { + // Initialize our stack if needed. + if i.stack == nil && i.node != nil { + i.stack = []rawStackEntry{ + rawStackEntry{ + edges: edges{ + edge{node: i.node}, + }, + }, + } + } + + for len(i.stack) > 0 { + // Inspect the last element of the stack. + n := len(i.stack) + last := i.stack[n-1] + elem := last.edges[0].node + + // Update the stack. + if len(last.edges) > 1 { + i.stack[n-1].edges = last.edges[1:] + } else { + i.stack = i.stack[:n-1] + } + + // Push the edges onto the frontier. + if len(elem.edges) > 0 { + path := last.path + string(elem.prefix) + i.stack = append(i.stack, rawStackEntry{path, elem.edges}) + } + + i.pos = elem + i.path = last.path + string(elem.prefix) + return + } + + i.pos = nil + i.path = "" +}