~themue/juju-core/053-env-more-script-friendly

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
package state

import (
	"fmt"
	"labix.org/v2/mgo"
	"labix.org/v2/mgo/txn"
	"launchpad.net/juju-core/utils"
	"strings"
)

// annotatorDoc represents the internal state of annotations for an Entity in
// MongoDB. Note that the annotations map is not maintained in local storage
// due to the fact that it is not accessed directly, but through
// Annotations/Annotation below.
// Note also the correspondence with AnnotationInfo in state/api/params.
type annotatorDoc struct {
	GlobalKey   string `bson:"_id"`
	Tag         string
	Annotations map[string]string
}

// annotator implements annotation-related methods
// for any entity that wishes to use it.
type annotator struct {
	globalKey string
	tag       string
	st        *State
}

// SetAnnotations adds key/value pairs to annotations in MongoDB.
func (a *annotator) SetAnnotations(pairs map[string]string) (err error) {
	defer utils.ErrorContextf(&err, "cannot update annotations on %s", a.tag)
	if len(pairs) == 0 {
		return nil
	}
	// Collect in separate maps pairs to be inserted/updated or removed.
	toRemove := make(map[string]bool)
	toInsert := make(map[string]string)
	toUpdate := make(map[string]string)
	for key, value := range pairs {
		if strings.Contains(key, ".") {
			return fmt.Errorf("invalid key %q", key)
		}
		if value == "" {
			toRemove["annotations."+key] = true
		} else {
			toInsert[key] = value
			toUpdate["annotations."+key] = value
		}
	}
	// Two attempts should be enough to update annotations even with racing
	// clients - if the document does not already exist, one of the clients
	// will create it and the others will fail, then all the rest of the
	// clients should succeed on their second attempt. If the referred-to
	// entity has disappeared, and removed its annotations in the meantime,
	// we consider that worthy of an error (will be fixed when new entities
	// can never share names with old ones).
	for i := 0; i < 2; i++ {
		var ops []txn.Op
		if count, err := a.st.annotations.FindId(a.globalKey).Count(); err != nil {
			return err
		} else if count == 0 {
			// Check that the annotator entity was not previously destroyed.
			if i != 0 {
				return fmt.Errorf("%s no longer exists", a.tag)
			}
			ops, err = a.insertOps(toInsert)
			if err != nil {
				return err
			}
		} else {
			ops = a.updateOps(toUpdate, toRemove)
		}
		if err := a.st.runner.Run(ops, "", nil); err == nil {
			return nil
		} else if err != txn.ErrAborted {
			return err
		}
	}
	return ErrExcessiveContention
}

// insertOps returns the operations required to insert annotations in MongoDB.
func (a *annotator) insertOps(toInsert map[string]string) ([]txn.Op, error) {
	tag := a.tag
	ops := []txn.Op{{
		C:      a.st.annotations.Name,
		Id:     a.globalKey,
		Assert: txn.DocMissing,
		Insert: &annotatorDoc{a.globalKey, tag, toInsert},
	}}
	if strings.HasPrefix(tag, "environment-") {
		return ops, nil
	}
	// If the entity is not the environment, add a DocExists check on the
	// entity document, in order to avoid possible races between entity
	// removal and annotation creation.
	coll, id, err := a.st.ParseTag(tag)
	if err != nil {
		return nil, err
	}
	return append(ops, txn.Op{
		C:      coll,
		Id:     id,
		Assert: txn.DocExists,
	}), nil
}

// updateOps returns the operations required to update or remove annotations in MongoDB.
func (a *annotator) updateOps(toUpdate map[string]string, toRemove map[string]bool) []txn.Op {
	return []txn.Op{{
		C:      a.st.annotations.Name,
		Id:     a.globalKey,
		Assert: txn.DocExists,
		Update: D{{"$set", toUpdate}, {"$unset", toRemove}},
	}}
}

// Annotations returns all the annotations corresponding to an entity.
func (a *annotator) Annotations() (map[string]string, error) {
	doc := new(annotatorDoc)
	err := a.st.annotations.FindId(a.globalKey).One(doc)
	if err == mgo.ErrNotFound {
		// Returning an empty map if there are no annotations.
		return make(map[string]string), nil
	}
	if err != nil {
		return nil, err
	}
	return doc.Annotations, nil
}

// Annotation returns the annotation value corresponding to the given key.
// If the requested annotation is not found, an empty string is returned.
func (a *annotator) Annotation(key string) (string, error) {
	ann, err := a.Annotations()
	if err != nil {
		return "", err
	}
	return ann[key], nil
}

// annotationRemoveOp returns an operation to remove a given annotation
// document from MongoDB.
func annotationRemoveOp(st *State, id string) txn.Op {
	return txn.Op{
		C:      st.annotations.Name,
		Id:     id,
		Remove: true,
	}
}