cmd/juju: working bootstrap and destroy commands

store/store.go

 // The store package is capable of storing and updating charms in a MongoDB
 // database, as well as maintaining further information about them such as
 // the VCS revision the charm was loaded from and the URLs for the charms.
 package store
 
 import (
         "crypto/sha256"
         "encoding/hex"
         "errors"
         "fmt"
         "hash"
         "io"
         "launchpad.net/juju/go/charm"
         "launchpad.net/juju/go/log"
         "launchpad.net/mgo"
         "launchpad.net/mgo/bson"
         "sort"
+        "strconv"
+        "sync"
         "time"
 )
 
 // The following MongoDB collections are currently used:
 //
 //     juju.events    - Log of events relating to the lifecycle of charms
 //     juju.charms    - Information about the stored charms
 //     juju.charmfs.* - GridFS with the charm files
 //     juju.locks     - Has unique keys with url of updating charms
 
 var (
         ErrUpdateConflict  = errors.New("charm update in progress")
         ErrRedundantUpdate = errors.New("charm is up-to-date")
         ErrNotFound        = errors.New("entry not found")
 )
 
 const (
         UpdateTimeout = 600e9
 )
 
 // Store holds a connection to a charm store.
 type Store struct {
         session *storeSession
+
+        // Cache for statistics key words (two generations).
+        cacheMu       sync.RWMutex
+        statsTokenNew map[string]int
+        statsTokenOld map[string]int
 }
 
 // Open creates a new session with the store. It connects to the MongoDB
 // server at the given address (as expected by the Mongo function in the
 // launchpad.net/mgo package).
 func Open(mongoAddr string) (store *Store, err error) {
         log.Printf("Store opened. Connecting to: %s", mongoAddr)
         store = &Store{}
         session, err := mgo.Dial(mongoAddr)
         if err != nil {
                 log.Printf("Error connecting to MongoDB: %v", err)
                 return nil, err
         }
-        store = &Store{&storeSession{session}}
+
+        store = &Store{session: &storeSession{session}}
+
+        // Ignore error. It'll always fail after created.
+        // TODO Check the error once mgo hands it to us.
+        _ = store.session.DB("juju").Run(bson.D{{"create", "stat.counters"}, {"autoIndexId", false}}, nil)
+
+        counters := store.session.StatCounters()
+        err = counters.EnsureIndex(mgo.Index{Key: []string{"k", "t"}, Unique: true})
+        if err != nil {
+                log.Printf("Error ensuring stat.counters index: %v", err)
+                session.Close()
+                return nil, err
+        }
+
         charms := store.session.Charms()
         err = charms.EnsureIndex(mgo.Index{Key: []string{"urls", "revision"}, Unique: true})
         if err != nil {
                 log.Printf("Error ensuring charms index: %v", err)
                 session.Close()
                 return nil, err
         }
+
         events := store.session.Events()
         err = events.EnsureIndex(mgo.Index{Key: []string{"urls", "digest"}})
         if err != nil {
                 log.Printf("Error ensuring events index: %v", err)
                 session.Close()
                 return nil, err
         }
-        // Put the socket we used on EnsureIndex back in the pool.
+
+        // Put the socket we used back in the pool.
         session.Refresh()
         return store, nil
 }
 
 // Close terminates the connection with the store.
 func (s *Store) Close() {
         s.session.Close()
+}
+
+// statsKey returns the compound statistics identifier that represents key.
+// If write is true, the identifier will be created if necessary.
+// Identifiers have a form similar to "ab:c:def:", where each section is a
+// base-32 number that represents the respective word in key. This form
+// allows efficiently indexing and searching for prefixes, while detaching
+// the key content and size from the actual words used in key.
+func (s *Store) statsKey(session *storeSession, key []string, write bool) (string, error) {
+        if len(key) == 0 {
+                return "", fmt.Errorf("store: empty statistics key")
+        }
+        tokens := session.StatTokens()
+        skey := make([]byte, 0, len(key)*4)
+        // Retry limit is mainly to prevent infinite recursion in edge cases,
+        // such as if the database is ever run in read-only mode.
+        // The logic below should deteministically stop in normal scenarios.
+        var err error
+        for i, retry := 0, 30; i < len(key) && retry > 0; retry-- {
+                id, found := s.statsTokenId(key[i])
+                if !found {
+                        var t struct {
+                                Id    int    "_id"
+                                Token string "t"
+                        }
+                        err = tokens.Find(bson.D{{"t", key[i]}}).One(&t)
+                        if err == mgo.NotFound {
+                                if !write {
+                                        return "", ErrNotFound
+                                }
+                                t.Id, err = tokens.Count()
+                                if err != nil {
+                                        continue
+                                }
+                                t.Id++
+                                t.Token = key[i]
+                                err = tokens.Insert(&t)
+                        }
+                        if err != nil {
+                                continue
+                        }
+                        s.cacheStatsTokenId(t.Token, t.Id)
+                        id = t.Id
+                }
+                skey = strconv.AppendInt(skey, int64(id), 32)
+                skey = append(skey, ':')
+                i++
+        }
+        if err != nil {
+                return "", err
+        }
+        return string(skey), nil
+}
+
+const statsTokenCacheSize = 512
+
+// cacheStatsTokenId adds the id for token into the cache.
+// The cache has two generations so that the least frequently used
+// tokens are evicted regularly.
+func (s *Store) cacheStatsTokenId(token string, id int) {
+        s.cacheMu.Lock()
+        defer s.cacheMu.Unlock()
+        // Can't possibly be >, but reviews want it for defensiveness.
+        if len(s.statsTokenNew) >= statsTokenCacheSize {
+                s.statsTokenOld = s.statsTokenNew
+                s.statsTokenNew = nil
+        }
+        if s.statsTokenNew == nil {
+                s.statsTokenNew = make(map[string]int, statsTokenCacheSize)
+        }
+        s.statsTokenNew[token] = id
+}
+
+// statsTokenId returns the id for token from the cache, if found.
+func (s *Store) statsTokenId(token string) (id int, found bool) {
+        s.cacheMu.RLock()
+        id, found = s.statsTokenNew[token]
+        if found {
+                s.cacheMu.RUnlock()
+                return
+        }
+        id, found = s.statsTokenOld[token]
+        s.cacheMu.RUnlock()
+        if found {
+                s.cacheStatsTokenId(token, id)
+        }
+        return
+}
+
+var counterEpoch = time.Date(2012, 1, 1, 0, 0, 0, 0, time.UTC).Unix()
+
+// IncCounter increases by one the counter associated with the composed key.
+func (s *Store) IncCounter(key []string) error {
+        session := s.session.Copy()
+        defer session.Close()
+
+        skey, err := s.statsKey(session, key, true)
+        if err != nil {
+                return err
+        }
+
+        t := time.Now().UTC()
+        // Round to the start of the minute so we get one document per minute at most.
+        t = t.Add(-time.Duration(t.Second()) * time.Second)
+        counters := session.StatCounters()
+        _, err = counters.Upsert(bson.D{{"k", skey}, {"t", int32(t.Unix() - counterEpoch)}}, bson.D{{"$inc", bson.D{{"c", 1}}}})
+        return err
+}
+
+// SumCounter returns the sum of all the counters that exactly match key,
+// or that are prefixed by it if prefix is true.
+func (s *Store) SumCounter(key []string, prefix bool) (count int64, err error) {
+        session := s.session.Copy()
+        defer session.Close()
+
+        skey, err := s.statsKey(session, key, false)
+        if err == ErrNotFound {
+                return 0, nil
+        }
+        if err != nil {
+                return 0, err
+        }
+
+        var regex string
+        if prefix {
+                regex = "^" + skey
+        } else {
+                regex = "^" + skey + "$"
+        }
+
+        job := mgo.MapReduce{
+                Map:    "function() { emit('count', this.c); }",
+                Reduce: "function(key, values) { return Array.sum(values); }",
+        }
+        var result []struct{ Value int64 }
+        counters := session.StatCounters()
+        _, err = counters.Find(bson.D{{"k", bson.D{{"$regex", regex}}}}).MapReduce(job, &result)
+        if len(result) > 0 {
+                return result[0].Value, err
+        }
+        return 0, err
 }
 
 // A CharmPublisher is responsible for importing a charm dir onto the store.
 type CharmPublisher struct {
         revision int
         w        *charmWriter
 }
 
 // Revision returns the revision that will be assigned to the published charm.
 func (p *CharmPublisher) Revision() int {
@@ skipping 131 matching lines @@
 }
 
 // finish completes the charm writing process and inserts the final metadata.
 // After it completes the charm will be available for consumption.
 func (w *charmWriter) finish() error {
         if w.file == nil {
                 return nil
         }
         defer w.session.Close()
         id := w.file.Id()
+        size := w.file.Size()
         err := w.file.Close()
         if err != nil {
                 log.Printf("Failed to close GridFS file: %v", err)
                 return err
         }
         charms := w.session.Charms()
         sha256 := hex.EncodeToString(w.sha256.Sum(nil))
         charm := charmDoc{
                 w.urls,
                 w.revision,
                 w.digest,
                 sha256,
+                size,
                 id.(bson.ObjectId),
                 w.charm.Meta(),
                 w.charm.Config(),
         }
         if err = charms.Insert(&charm); err != nil {
                 err = maybeConflict(err)
                 log.Printf("Failed to insert new revision of charm %v: %v", w.urls, err)
                 return err
         }
         return nil
 }
 
 type CharmInfo struct {
         revision int
         digest   string
         sha256   string
+        size     int64
         fileId   bson.ObjectId
         meta     *charm.Meta
         config   *charm.Config
 }
 
 // Statically ensure CharmInfo is a charm.Charm.
 var _ charm.Charm = (*CharmInfo)(nil)
 
 // Revision returns the store charm's revision.
 func (ci *CharmInfo) Revision() int {
         return ci.revision
 }
 
 // BundleSha256 returns the sha256 checksum for the stored charm bundle.
 func (ci *CharmInfo) BundleSha256() string {
         return ci.sha256
+}
+
+// BundleSize returns the size for the stored charm bundle.
+func (ci *CharmInfo) BundleSize() int64 {
+        return ci.size
 }
 
 // Digest returns the unique identifier that represents the charm
 // data imported. This is typically set to the VCS revision digest.
 func (ci *CharmInfo) Digest() string {
         return ci.digest
 }
 
 // Meta returns the charm.Meta details for the stored charm.
 func (ci *CharmInfo) Meta() *charm.Meta {
@@ skipping 20 matching lines @@
         if rev == -1 {
                 qdoc = bson.D{{"urls", url}}
         } else {
                 qdoc = bson.D{{"urls", url}, {"revision", rev}}
         }
         err = charms.Find(qdoc).Sort(bson.D{{"revision", -1}}).One(&cdoc)
         if err != nil {
                 log.Printf("Failed to find charm %s: %v", url, err)
                 return nil, ErrNotFound
         }
-        return &CharmInfo{cdoc.Revision, cdoc.Digest, cdoc.Sha256, cdoc.FileId, cdoc.Meta, cdoc.Config}, nil
+        info = &CharmInfo{
+                cdoc.Revision,
+                cdoc.Digest,
+                cdoc.Sha256,
+                cdoc.Size,
+                cdoc.FileId,
+                cdoc.Meta,
+                cdoc.Config,
+        }
+        return info, nil
 }
 
 // OpenCharm opens for reading via rc the charm currently available at url.
 // rc must be closed after dealing with it or resources will leak.
 func (s *Store) OpenCharm(url *charm.URL) (info *CharmInfo, rc io.ReadCloser, err error) {
         log.Debugf("Opening charm %s", url)
         info, err = s.CharmInfo(url)
         if err != nil {
                 return nil, nil, err
         }
@@ skipping 24 matching lines @@
         r.session.Close()
         return err
 }
 
 // charmDoc represents the document stored in MongoDB for a charm.
 type charmDoc struct {
         URLs     []*charm.URL
         Revision int
         Digest   string
         Sha256   string
+        Size     int64
         FileId   bson.ObjectId
         Meta     *charm.Meta
         Config   *charm.Config
 }
 
 // LockUpdates acquires a server-side lock for updating a single charm
 // that is supposed to be made available in all of the provided urls.
 // If the lock can't be acquired in any of the urls, an error will be
 // immediately returned.
 // In the usual case, any locking done is undone when an error happens,
@@ skipping 107 matching lines @@
 // Events returns the mongo collection where charm events are stored.
 func (s *storeSession) Events() *mgo.Collection {
         return s.DB("juju").C("events")
 }
 
 // Locks returns the mongo collection where charm locks are stored.
 func (s *storeSession) Locks() *mgo.Collection {
         return s.DB("juju").C("locks")
 }
 
+// StatTokens returns the mongo collection for storing key tokens
+// for statistics collection.
+func (s *storeSession) StatTokens() *mgo.Collection {
+        return s.DB("juju").C("stat.tokens")
+}
+
+// StatCounters returns the mongo collection for counter values.
+func (s *storeSession) StatCounters() *mgo.Collection {
+        return s.DB("juju").C("stat.counters")
+}
+
 type CharmEventKind int
 
 const (
         EventPublished CharmEventKind = iota + 1
         EventPublishError
 
         EventKindCount
 )
 
 func (k CharmEventKind) String() string {
@@ skipping 66 matching lines @@
 func mustLackRevision(context string, urls ...*charm.URL) error {
         for _, url := range urls {
                 if url.Revision != -1 {
                         err := fmt.Errorf("%s: got charm URL with revision: %s", context, url)
                         log.Printf("%v", err)
                         return err
                 }
         }
         return nil
 }