github.com/axw/juju@v0.0.0-20161005053422-4bd6544d08d4/state/minimumunits.go

// Copyright 2012, 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package state

import (
	"github.com/juju/errors"
	jujutxn "github.com/juju/txn"
	"gopkg.in/mgo.v2/bson"
	"gopkg.in/mgo.v2/txn"
)

// minUnitsDoc keeps track of relevant changes on the service's MinUnits field
// and on the number of alive units for the application.
// A new document is created when MinUnits is set to a non zero value.
// A document is deleted when either the associated service is destroyed
// or MinUnits is restored to zero. The Revno is increased when either MinUnits
// for a service is increased or a unit is destroyed.
// TODO(frankban): the MinUnitsWatcher reacts to changes by sending events,
// each one describing one or more services. A worker reacts to those events
// ensuring the number of units for the application is never less than the actual
// alive units: new units are added if required.
type minUnitsDoc struct {
	// ApplicationName is safe to be used here in place of its globalKey, since
	// the referred entity type is always the Service.
	DocID           string `bson:"_id"`
	ApplicationName string
	ModelUUID       string `bson:"model-uuid"`
	Revno           int
}

// SetMinUnits changes the number of minimum units required by the service.
func (s *Application) SetMinUnits(minUnits int) (err error) {
	defer errors.DeferredAnnotatef(&err, "cannot set minimum units for application %q", s)
	defer func() {
		if err == nil {
			s.doc.MinUnits = minUnits
		}
	}()
	if minUnits < 0 {
		return errors.New("cannot set a negative minimum number of units")
	}
	service := &Application{st: s.st, doc: s.doc}
	// Removing the document never fails. Racing clients trying to create the
	// document generate one failure, but the second attempt should succeed.
	// If one client tries to update the document, and a racing client removes
	// it, the former should be able to re-create the document in the second
	// attempt. If the referred-to service advanced its life cycle to a not
	// alive state, an error is returned after the first failing attempt.
	buildTxn := func(attempt int) ([]txn.Op, error) {
		if attempt > 0 {
			if err := service.Refresh(); err != nil {
				return nil, err
			}
		}
		if service.doc.Life != Alive {
			return nil, errors.New("application is no longer alive")
		}
		if minUnits == service.doc.MinUnits {
			return nil, jujutxn.ErrNoOperations
		}
		return setMinUnitsOps(service, minUnits), nil
	}
	return s.st.run(buildTxn)
}

// setMinUnitsOps returns the operations required to set MinUnits on the
// service and to create/update/remove the minUnits document in MongoDB.
func setMinUnitsOps(service *Application, minUnits int) []txn.Op {
	state := service.st
	applicationname := service.Name()
	ops := []txn.Op{{
		C:      applicationsC,
		Id:     state.docID(applicationname),
		Assert: isAliveDoc,
		Update: bson.D{{"$set", bson.D{{"minunits", minUnits}}}},
	}}
	if service.doc.MinUnits == 0 {
		return append(ops, txn.Op{
			C:      minUnitsC,
			Id:     state.docID(applicationname),
			Assert: txn.DocMissing,
			Insert: &minUnitsDoc{
				ApplicationName: applicationname,
				ModelUUID:       service.st.ModelUUID(),
			},
		})
	}
	if minUnits == 0 {
		return append(ops, minUnitsRemoveOp(state, applicationname))
	}
	if minUnits > service.doc.MinUnits {
		op := minUnitsTriggerOp(state, applicationname)
		op.Assert = txn.DocExists
		return append(ops, op)
	}
	return ops
}

// minUnitsTriggerOp returns the operation required to increase the minimum
// units revno for the application in MongoDB, ignoring the case of document not
// existing. This is included in the operations performed when a unit is
// destroyed: if the document exists, then we need to update the Revno.
// If the service does not require a minimum number of units, then the
// operation is a noop.
func minUnitsTriggerOp(st *State, applicationname string) txn.Op {
	return txn.Op{
		C:      minUnitsC,
		Id:     st.docID(applicationname),
		Update: bson.D{{"$inc", bson.D{{"revno", 1}}}},
	}
}

// minUnitsRemoveOp returns the operation required to remove the minimum
// units document from MongoDB.
func minUnitsRemoveOp(st *State, applicationname string) txn.Op {
	return txn.Op{
		C:      minUnitsC,
		Id:     st.docID(applicationname),
		Remove: true,
	}
}

// MinUnits returns the minimum units count for the application.
func (s *Application) MinUnits() int {
	return s.doc.MinUnits
}

// EnsureMinUnits adds new units if the service's MinUnits value is greater
// than the number of alive units.
func (s *Application) EnsureMinUnits() (err error) {
	defer errors.DeferredAnnotatef(&err, "cannot ensure minimum units for application %q", s)
	service := &Application{st: s.st, doc: s.doc}
	for {
		// Ensure the service is alive.
		if service.doc.Life != Alive {
			return errors.New("application is not alive")
		}
		// Exit without errors if the MinUnits for the application is not set.
		if service.doc.MinUnits == 0 {
			return nil
		}
		// Retrieve the number of alive units for the application.
		aliveUnits, err := aliveUnitsCount(service)
		if err != nil {
			return err
		}
		// Calculate the number of required units to be added.
		missing := service.doc.MinUnits - aliveUnits
		if missing <= 0 {
			return nil
		}
		name, ops, err := ensureMinUnitsOps(service)
		if err != nil {
			return err
		}
		// Add missing unit.
		switch err := s.st.runTransaction(ops); err {
		case nil:
			// Assign the new unit.
			unit, err := s.st.Unit(name)
			if err != nil {
				return err
			}
			if err := service.st.AssignUnit(unit, AssignNew); err != nil {
				return err
			}
			// No need to proceed and refresh the service if this was the
			// last/only missing unit.
			if missing == 1 {
				return nil
			}
		case txn.ErrAborted:
			// Refresh the service and restart the loop.
		default:
			return err
		}
		if err := service.Refresh(); err != nil {
			return err
		}
	}
}

// aliveUnitsCount returns the number a alive units for the application.
func aliveUnitsCount(service *Application) (int, error) {
	units, closer := service.st.getCollection(unitsC)
	defer closer()

	query := bson.D{{"application", service.doc.Name}, {"life", Alive}}
	return units.Find(query).Count()
}

// ensureMinUnitsOps returns the operations required to add a unit for the
// service in MongoDB and the name for the new unit. The resulting transaction
// will be aborted if the service document changes when running the operations.
func ensureMinUnitsOps(service *Application) (string, []txn.Op, error) {
	asserts := bson.D{{"txn-revno", service.doc.TxnRevno}}
	return service.addUnitOps("", asserts)
}