func (s *DatamodelSuite) TestReserveRouterPortAndUpdateTrie(c *C) {
Zk.RecursiveDelete(helper.GetBaseRouterPortsPath(true))
Zk.RecursiveDelete(helper.GetBaseRouterPortsPath(false))
Zk.RecursiveDelete(helper.GetBaseLockPath())
Zk.RecursiveDelete("/atlantis/router")
CreateRouterPaths()
CreateRouterPortsPaths()
CreateLockPaths()
MinRouterPort = uint16(65533)
MaxRouterPort = uint16(65535)
helper.SetRouterRoot(true)
port, created, err := ReserveRouterPortAndUpdateTrie(true, "app", "sha", "env")
c.Assert(err, IsNil)
c.Assert(created, Equals, true)
c.Assert(port, Equals, "65533")
trie, err := routerzk.GetTrie(Zk.Conn, helper.GetAppEnvTrieName("app", "env"))
c.Assert(err, IsNil)
c.Assert(len(trie.Rules), Equals, 1)
port, created, err = ReserveRouterPortAndUpdateTrie(true, "app", "sha2", "env")
c.Assert(err, IsNil)
c.Assert(created, Equals, false)
c.Assert(port, Equals, "65533")
trie, err = routerzk.GetTrie(Zk.Conn, helper.GetAppEnvTrieName("app", "env"))
c.Assert(err, IsNil)
c.Assert(len(trie.Rules), Equals, 2)
}
func GetTrie(name string) (cfg.Trie, error) {
if name == "" {
return cfg.Trie{}, errors.New("Please specify a name")
}
return routerzk.GetTrie(zkConn.Conn, name)
}
func CleanupCreatedPoolRefs(internal bool, app, sha, env string) error {
helper.SetRouterRoot(internal)
// remove static rule, cleanup rule from trie if needed
ruleName := helper.GetAppShaEnvStaticRuleName(app, sha, env)
trieName := helper.GetAppEnvTrieName(app, env)
// remove static rule from trie
trie, err := routerzk.GetTrie(Zk.Conn, trieName)
if err != nil {
return err
}
newRules := []string{}
for _, rule := range trie.Rules {
if rule != ruleName {
newRules = append(newRules, rule)
}
}
if len(trie.Rules) != len(newRules) {
trie.Rules = newRules
err = routerzk.SetTrie(Zk.Conn, trie)
if err != nil {
return err
}
}
// delete static rule
err = routerzk.DelRule(Zk.Conn, ruleName)
if err != nil {
return err
}
return nil
}
func ParseTrie(name string) (Trie, error) {
var res Trie
data, err := zk.GetTrie(Zk, name)
if err != nil {
return res, err
}
return Trie{data.Rules}, nil
}
func (e *GetTrieExecutor) Execute(t *Task) (err error) {
if e.arg.Name == "" {
return errors.New("Please specify a name")
}
helper.SetRouterRoot(e.arg.Internal)
e.reply.Trie, err = routerzk.GetTrie(datamodel.Zk.Conn, e.arg.Name)
if err != nil {
e.reply.Status = StatusError
} else {
e.reply.Status = StatusOk
}
return err
}
func UpdateAppEnvTrie(internal bool, app, sha, env string) (string, error) {
helper.SetRouterRoot(internal)
// create trie (if it doesn't exist)
trieName := helper.GetAppEnvTrieName(app, env)
if exists, err := routerzk.TrieExists(Zk.Conn, trieName); !exists || err != nil {
err = routerzk.SetTrie(Zk.Conn, routercfg.Trie{
Name: trieName,
Rules: []string{},
Internal: internal,
})
if err != nil {
return trieName, err
}
}
// if sha != "" attach pool as static rule (if trie is empty)
if sha != "" {
// if static rule does not exist, create it
ruleName := helper.GetAppShaEnvStaticRuleName(app, sha, env)
poolName := helper.CreatePoolName(app, sha, env)
if exists, err := routerzk.RuleExists(Zk.Conn, ruleName); !exists || err != nil {
err = routerzk.SetRule(Zk.Conn, routercfg.Rule{
Name: ruleName,
Type: "static",
Value: "true",
Pool: poolName,
Internal: internal,
})
if err != nil {
return trieName, err
}
}
trie, err := routerzk.GetTrie(Zk.Conn, trieName)
if err != nil {
return trieName, err
}
if len(trie.Rules) == 0 {
trie.Rules = []string{ruleName}
} else {
trie.Rules = append(trie.Rules, ruleName)
}
if err = routerzk.SetTrie(Zk.Conn, trie); err != nil {
return trieName, err
}
}
return trieName, nil
}