// NewAgent returns an initalized Agent.
func NewAgent(prefix string, c *consul.Client, d *docker.Client) *Agent {
return &Agent{
KV: c.KV(),
Docker: d,
Prefix: prefix,
}
}
func writeToConsul(t *testing.T, prefix, key string, client *consulapi.Client) []byte {
token := os.Getenv("TOKEN")
dc := os.Getenv("DC")
if dc == "" {
dc = "dc1"
}
kv := client.KV()
writeOptions := &consulapi.WriteOptions{Token: token, Datacenter: dc}
// Delete all keys in the prefixed KV space
if _, err := kv.DeleteTree(prefix, writeOptions); err != nil {
t.Fatalf("err: %v", err)
}
// Put a test KV
encodedValue := make([]byte, base64.StdEncoding.EncodedLen(1024))
base64.StdEncoding.Encode(encodedValue, createRandomBytes(1024))
p := &consulapi.KVPair{Key: key, Flags: 42, Value: encodedValue}
if _, err := kv.Put(p, writeOptions); err != nil {
t.Fatalf("err: %v", err)
}
return encodedValue
}
func writeFileToConsul(t *testing.T, prefix, key string, file string, client *consulapi.Client) []byte {
token := os.Getenv("TOKEN")
dc := os.Getenv("DC")
if dc == "" {
dc = "dc1"
}
kv := client.KV()
writeOptions := &consulapi.WriteOptions{Token: token, Datacenter: dc}
// Delete all keys in the prefixed KV space
if _, err := kv.DeleteTree(prefix, writeOptions); err != nil {
t.Fatalf("err: %v", err)
}
fileBytes, err := ioutil.ReadFile(file)
if err != nil {
t.Fatalf("err: %v", err)
}
p := &consulapi.KVPair{Key: key, Flags: 42, Value: fileBytes}
if _, err := kv.Put(p, writeOptions); err != nil {
t.Fatalf("err: %v", err)
}
return fileBytes
}
func NewInstall(consulClient *consul.Client, marathon *marathon.Marathon, mesos *mesos.Mesos, zkHosts []string) (*Install, error) {
apiConfig = map[string]interface{}{
"mantl": map[string]interface{}{
"zookeeper": map[string]interface{}{
"hosts": strings.Join(zkHosts, ","),
},
},
}
if mesos != nil {
mesosAuthRequired, err := mesos.RequiresAuthentication()
if err != nil {
return nil, err
}
mesosConfig := map[string]interface{}{
"principal": mesos.Principal,
"secret": mesos.Secret,
"secret-path": mesos.SecretPath,
"authentication-enabled": mesosAuthRequired,
}
mantlConfig := apiConfig["mantl"].(map[string]interface{})
mantlConfig["mesos"] = mesosConfig
}
zookeeper := zookeeper.NewZookeeper(zkHosts)
return &Install{consulClient, consulClient.KV(), marathon, mesos, zookeeper}, nil
}
func getAttributes(keys []string, overwriteAttributes map[string]interface{}) (map[string]interface{}, error) {
var attributes map[string]interface{}
var c *api.Client = util.Consul()
attributes = make(map[string]interface{})
// Get attributes from consul KVS
for _, key := range keys {
list, _, err := c.KV().List(key, &api.QueryOptions{})
if err != nil {
return nil, err
}
for _, kv := range list {
var a map[string]interface{}
if err := json.Unmarshal(kv.Value, &a); err != nil {
return nil, err
}
if err := mergo.MergeWithOverwrite(&attributes, a); err != nil {
return nil, errors.New(fmt.Sprintf("Failed to merge attributes(%v)", err))
}
}
}
// Overwrite some attributes by specified parameter in task.yml
if err := mergo.MergeWithOverwrite(&attributes, overwriteAttributes); err != nil {
return nil, err
}
return attributes, nil
}
// NewWaiter creates a new Wait entry with a sensible default isReady function if
// nil is provided in its place.
func NewWaiter(client *api.Client, prefix string, minimumNodes int, isReady func(n *WaitNode) bool) *Wait {
if isReady == nil {
isReady = func(n *WaitNode) bool {
return true
}
}
nodeUpdateCh := make(chan WaitNodeUpdate, 2)
nodeReadyCh := make(chan WaitNode, 2)
allReadyCh := make(chan []WaitNode, 2)
return &Wait{
Prefix: prefix,
MinimumNodes: minimumNodes,
IsReady: isReady,
NodeUpdate: nodeUpdateCh,
NodeReady: nodeReadyCh,
AllReady: allReadyCh,
nodeUpdate: nodeUpdateCh,
nodeReady: nodeReadyCh,
allReady: allReadyCh,
client: client,
kv: client.KV(),
}
}
/*
Get full tree under a key and optionally unmarshal.
Args:
client : Consul client
key : Key to query for.
output : Unmarshal data to this interface{} if non-nil
*/
func GetKVTree(client *api.Client, key string, output interface{}) (pairs api.KVPairs, err error) {
kv := client.KV()
pairs, _, err = kv.List(key, nil)
if output != nil {
err = Unmarshal(pairs, output)
}
return
}
func putKV(client *api.Client, key, value string, index uint64) (bool, error) {
p := &api.KVPair{Key: key[1:], Value: []byte(value), ModifyIndex: index}
ok, _, err := client.KV().CAS(p, nil)
if err != nil {
return false, err
}
return ok, nil
}
// consulDel removes a key from the Consul KV store.
func consulDel(c *consul.Client, key string) (bool, error) {
kv := c.KV()
key = strings.TrimPrefix(key, "/")
_, err := kv.Delete(key, nil)
if err != nil {
return false, err
}
Log(fmt.Sprintf("action='consulDel' key='%s'", key), "info")
return true, err
}
// Set sets a key's value inside Consul.
func Set(c *consul.Client, key, value string) bool {
p := &consul.KVPair{Key: key, Value: []byte(value)}
kv := c.KV()
Log(fmt.Sprintf("key='%s' value='%s'", key, value), "info")
_, err := kv.Put(p, nil)
if err != nil {
panic(err)
}
return true
}
// consulSet a value for a key in the Consul KV store.
func consulSet(c *consul.Client, key string, value string) (bool, error) {
key = strings.TrimPrefix(key, "/")
p := &consul.KVPair{Key: key, Value: []byte(value)}
kv := c.KV()
_, err := kv.Put(p, nil)
if err != nil {
return false, err
}
Log(fmt.Sprintf("action='consulSet' key='%s'", key), "debug")
return true, err
}
func putKey(c *consulapi.Client, k string, v string, s *Seeder) error {
writeOpt := &consulapi.WriteOptions{}
var err error
writeKey := &consulapi.KVPair{Key: k, Value: []byte(v)}
s.Data = append(s.Data, writeKey)
_, err = c.KV().Put(writeKey, writeOpt)
if err != nil {
return fmt.Errorf("Failed to write test key: %v", err)
}
return nil
}
func getKV(client *api.Client, key string, waitIndex uint64) (string, uint64, error) {
q := &api.QueryOptions{RequireConsistent: true, WaitIndex: waitIndex}
kvpair, meta, err := client.KV().Get(key, q)
if err != nil {
return "", 0, err
}
if kvpair == nil {
return "", meta.LastIndex, nil
}
return strings.TrimSpace(string(kvpair.Value)), meta.LastIndex, nil
}
func watch(
client *consulapi.Client,
prefix string,
path string,
token string,
pairCh chan<- consulapi.KVPairs,
errCh chan<- error,
quitCh <-chan struct{}) {
// Create the root for KVs, if necessary
mkdirp.Mk(path, 0777)
// Get the initial list of k/v pairs. We don't do a retryableList
// here because we want a fast fail if the initial request fails.
opts := &consulapi.QueryOptions{Token: token}
pairs, meta, err := client.KV().List(prefix, opts)
if err != nil {
errCh <- err
return
}
// Send the initial list out right away
pairCh <- pairs
// Loop forever (or until quitCh is closed) and watch the keys
// for changes.
curIndex := meta.LastIndex
for {
select {
case <-quitCh:
return
default:
}
pairs, meta, err = retryableList(
func() (consulapi.KVPairs, *consulapi.QueryMeta, error) {
opts = &consulapi.QueryOptions{WaitIndex: curIndex, Token: token}
return client.KV().List(prefix, opts)
})
if err != nil {
// This happens when the connection to the consul agent dies. Build in a retry by looping after a delay.
log.Warn("Error communicating with consul agent.")
continue
}
pairCh <- pairs
log.WithFields(log.Fields{
"curIndex": curIndex,
"lastIndex": meta.LastIndex,
}).Debug("Potential index update observed")
curIndex = meta.LastIndex
}
}
func checkKey(c *consulapi.Client, kv *consulapi.KVPair) error {
queryOpt := &consulapi.QueryOptions{}
keyCheck, _, err := c.KV().Get(kv.Key, queryOpt)
if err != nil || keyCheck == nil {
return fmt.Errorf("Failed to get key: %v", err)
}
reflecttest := reflect.DeepEqual(keyCheck.Value, kv.Value)
if reflecttest != true {
return fmt.Errorf("Key %v did not match\n\tExpected: %v\n\tGot: %v", kv.Key, string(kv.Value), string(keyCheck.Value))
}
return nil
}
// Get grabs a key's value from Consul.
func Get(c *consul.Client, key string) string {
var value string
kv := c.KV()
pair, _, err := kv.Get(key, nil)
if err != nil {
panic(err)
} else {
if pair != nil {
value = string(pair.Value[:])
} else {
value = ""
}
}
return value
}
func deleteKeyFromConsul(t *testing.T, key string, client *consulapi.Client) {
token := os.Getenv("TOKEN")
dc := os.Getenv("DC")
if dc == "" {
dc = "dc1"
}
kv := client.KV()
writeOptions := &consulapi.WriteOptions{Token: token, Datacenter: dc}
if _, err := kv.Delete(key, writeOptions); err != nil {
t.Fatalf("err: %v", err)
}
}
func getCerts(client *api.Client, key string, waitIndex uint64) (pemBlocks map[string][]byte, lastIndex uint64, err error) {
q := &api.QueryOptions{RequireConsistent: true, WaitIndex: waitIndex}
kvpairs, meta, err := client.KV().List(key, q)
if err != nil {
return nil, 0, fmt.Errorf("consul: list: %s", err)
}
if len(kvpairs) == 0 {
return nil, meta.LastIndex, nil
}
pemBlocks = map[string][]byte{}
for _, kvpair := range kvpairs {
pemBlocks[path.Base(kvpair.Key)] = kvpair.Value
}
return pemBlocks, meta.LastIndex, nil
}
// consulGet the value from a key in the Consul KV store.
func consulGet(c *consul.Client, key string) (string, error) {
var value string
kv := c.KV()
key = strings.TrimPrefix(key, "/")
pair, _, err := kv.Get(key, nil)
if err != nil {
return "", err
}
if pair != nil {
value = string(pair.Value[:])
} else {
value = ""
}
Log(fmt.Sprintf("action='consulGet' key='%s'", key), "debug")
return value, err
}
func (o *Operation) runRollout(client *api.Client) error {
kv := client.KV()
_, err := kv.Put(&api.KVPair{
Key: fmt.Sprintf("%s/current", strings.Trim(o.Config.Prefix, "/")),
Value: []byte(o.Version),
}, nil)
if err != nil {
o.UI.Error(fmt.Sprintf("Version '%s' could not be marked for rollout: %s", o.Version, err))
return err
}
o.UI.Info(fmt.Sprintf("Version '%s' marked for rollout", o.Version))
return nil
}
func (config *Config) fillFromConsul(client *consul.Client, appName string) error {
kv := client.KV()
structType := reflect.TypeOf(*config)
structValue := reflect.ValueOf(config).Elem()
for i := 0; i < structType.NumField(); i++ {
field := structType.Field(i)
consulKey := field.Tag.Get("consul")
pair, _, err := kv.Get(appName+"/"+consulKey, nil)
if pair == nil {
return err
}
consulValue := string(pair.Value)
structValue.FieldByName(field.Name).SetString(consulValue)
}
return nil
}
func NewInstall(consulClient *consul.Client, marathon *marathon.Marathon, mesos *mesos.Mesos, zookeeper *zookeeper.Zookeeper) (*Install, error) {
if mesos != nil {
mesosAuthRequired, err := mesos.RequiresAuthentication()
if err != nil {
return nil, err
}
apiConfig = map[string]interface{}{
"mantl": map[string]interface{}{
"mesos": map[string]interface{}{
"principal": mesos.Principal,
"secret": mesos.Secret,
"authentication-enabled": mesosAuthRequired,
},
},
}
}
return &Install{consulClient, consulClient.KV(), marathon, mesos, zookeeper}, nil
}
func testClientWrite(t *testing.T, client *consul.Client, data map[string]string) func() {
for k, v := range data {
_, err := client.KV().Put(&consul.KVPair{
Key: k,
Value: []byte(v),
}, nil)
if err != nil {
t.Fatalf("error writing to Consul: %s", err)
}
}
return func() {
for k, _ := range data {
_, err := client.KV().Delete(k, nil)
if err != nil {
t.Fatalf("error deleting from Consul: %s", err)
}
}
}
}
func NewConsulClient(addr string, dialTimeout time.Duration) (*ConsulClient, error) {
var c *api.Client
// configure address (and scheme if necessary)
cfg := api.DefaultConfig()
cfg.Address = addr
cfg.HttpClient = &http.Client{
Transport: &http.Transport{
Dial: func(network, addr string) (net.Conn, error) {
return net.DialTimeout(network, addr, dialTimeout)
},
},
}
// create custom client
c, err := api.NewClient(cfg)
if err != nil {
return nil, err
}
return &ConsulClient{addr: addr, client: c, kv: c.KV()}, nil
}
func (w *awaitValue) wait(client *api.Client, timeout time.Duration) error {
kv := client.KV()
startTime := time.Now()
for time.Now().Sub(startTime) < timeout {
p, _, err := kv.Get(w.Key, nil)
if err != nil {
return err
}
if p != nil && string(p.Value) == w.Value {
return nil
}
if p == nil && w.Value == "" {
return nil
}
}
return fmt.Errorf("Timeout expired waiting for '%s'='%s'", w.Key, w.Value)
}
// LoadConsulOverrides loads configuration from Consul (if any) and merges with current configuration
func LoadConsulOverrides(client *api.Client, prefix string) (map[string]interface{}, error) {
kv := client.KV()
prefix = strings.TrimPrefix(prefix, "/")
var (
pairs api.KVPairs
err error
)
for attempt := 0; attempt < 180; attempt++ {
pairs, _, err = kv.List(prefix, nil)
if err == nil {
break
}
log.Printf("Failure loading config from Consul: %s, retrying in 1 second, attempt %d", err, attempt)
time.Sleep(time.Second)
}
if err != nil {
return nil, fmt.Errorf("failure listing keys under %#v: %s", prefix, err)
}
configParts := map[string]interface{}{}
for _, pair := range pairs {
var val interface{}
err = json.Unmarshal(pair.Value, &val)
if err != nil {
return nil, fmt.Errorf("unable to decode JSON for key %#v: %s", pair.Key, err)
}
configParts[pair.Key] = val
}
return configParts, nil
}
// acquireKey tries to acquire a consul leader key. If successful we attain mastership.
func acquireKey(cl *api.Client, key string, ttl int, sessionName string) (string, *sequencer, bool, error) {
session := cl.Session()
entry := &api.SessionEntry{
TTL: strconv.Itoa(ttl) + "s", // ttl in seconds
Name: sessionName,
Behavior: api.SessionBehaviorDelete,
}
id, _, err := session.Create(entry, nil)
if err != nil {
log.Printf("Error while creating session: %v", err)
return "", nil, false, err
}
// Get a handle to the KV API
kv := cl.KV()
val := getHostname() + ":" + strconv.Itoa(os.Getpid()) // set value in 'hostname:pid' format
// PUT a new KV pair
p := &api.KVPair{
Key: key,
Session: id,
Value: []byte(val),
}
success, _, err := kv.Acquire(p, nil)
if err != nil {
log.Printf("Error while aquiring key: %v", err)
return "", nil, false, err
}
// get the sequencer
seq, err := getSequencer(kv, key)
if err != nil {
log.Printf("Error while retrieving sequencer %v", err)
return "", nil, false, err
}
return id, seq, success, nil
}
func (d *DedupManager) watchTemplate(client *consulapi.Client, t *Template) {
log.Printf("[INFO] (dedup) starting watch for template hash %s", t.hexMD5)
path := path.Join(d.config.Deduplicate.Prefix, t.hexMD5, "data")
// Determine if stale queries are allowed
var allowStale bool
if d.config.MaxStale != 0 {
allowStale = true
}
// Setup our query options
opts := &consulapi.QueryOptions{
AllowStale: allowStale,
WaitTime: 60 * time.Second,
}
START:
// Stop listening if we're stopped
select {
case <-d.stopCh:
return
default:
}
// If we are current the leader, wait for leadership lost
d.leaderLock.RLock()
lockCh, ok := d.leader[t]
d.leaderLock.RUnlock()
if ok {
select {
case <-lockCh:
goto START
case <-d.stopCh:
return
}
}
// Block for updates on the data key
log.Printf("[INFO] (dedup) listing data for template hash %s", t.hexMD5)
pair, meta, err := client.KV().Get(path, opts)
if err != nil {
log.Printf("[ERR] (dedup) failed to get '%s': %v", path, err)
select {
case <-time.After(listRetry):
goto START
case <-d.stopCh:
return
}
}
opts.WaitIndex = meta.LastIndex
// If we've exceeded the maximum staleness, retry without stale
if allowStale && meta.LastContact > d.config.MaxStale {
allowStale = false
log.Printf("[DEBUG] (dedup) %s stale data (last contact exceeded max_stale)", path)
goto START
}
// Re-enable stale queries if allowed
if d.config.MaxStale != 0 {
allowStale = true
}
// Stop listening if we're stopped
select {
case <-d.stopCh:
return
default:
}
// If we are current the leader, wait for leadership lost
d.leaderLock.RLock()
lockCh, ok = d.leader[t]
d.leaderLock.RUnlock()
if ok {
select {
case <-lockCh:
goto START
case <-d.stopCh:
return
}
}
// Parse the data file
if pair != nil && pair.Flags == templateDataFlag {
d.parseData(pair.Key, pair.Value)
}
goto START
}
func testConsul(client *consul.Client) error {
kv := client.KV()
_, _, err := kv.Get("mantl-install", nil)
return err
}
func NewSessionNoChecks(sessionName string, ttl time.Duration, client *api.Client, sessionMgr SessionManager) (*Session, error) {
return newSession(sessionName, ttl, true, client.KV(), sessionMgr)
}