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

package state

import (
	"fmt"

	"github.com/juju/errors"
	"github.com/juju/juju/mongo"

	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
)

type sequenceDoc struct {
	DocID     string `bson:"_id"`
	Name      string `bson:"name"`
	ModelUUID string `bson:"model-uuid"`
	Counter   int    // This stores the *next* value to return.
}

// sequence safely increments a database backed sequence, returning
// the next value.
func (st *State) sequence(name string) (int, error) {
	sequences, closer := st.getCollection(sequenceC)
	defer closer()
	query := sequences.FindId(name)
	inc := mgo.Change{
		Update: bson.M{
			"$set": bson.M{
				"name":       name,
				"model-uuid": st.ModelUUID(),
			},
			"$inc": bson.M{"counter": 1},
		},
		Upsert: true,
	}
	result := &sequenceDoc{}
	_, err := query.Apply(inc, result)
	if err != nil {
		return -1, fmt.Errorf("cannot increment %q sequence number: %v", name, err)
	}
	return result.Counter, nil
}

// sequenceWithMin safely increments a database backed sequence,
// allowing for a minimum value for the sequence to be specified. The
// minimum value is used as an initial value for the first use of a
// particular sequence. The minimum value will also cause a sequence
// value to jump ahead if the minimum is provided that is higher than
// the current sequence value.
//
// The data manipulated by `sequence` and `sequenceWithMin` is the
// same. It is safe to mix the 2 methods for the same sequence.
//
// `sequence` is more efficient than `sequenceWithMin` and should be
// preferred if there is no minimum value requirement.
func (st *State) sequenceWithMin(name string, minVal int) (int, error) {
	sequences, closer := st.getCollection(sequenceC)
	defer closer()
	updater := &dbSeqUpdater{
		coll: sequences,
		id:   st.docID(name),
		name: name,
	}
	return updateSeqWithMin(updater, minVal)
}

// seqUpdater abstracts away the database operations required for
// updating a sequence.
type seqUpdater interface {
	// read returns the current value of the sequence. If the sequence
	// doesn't exist yet it returns 0.
	read() (int, error)

	// create attempts to create a new sequence with the initial value
	// provided. It returns (true, nil) on success, (false, nil) if
	// the sequence already existed and (false, <some error>) if any
	// other error occurred.
	create(value int) (bool, error)

	// set attempts to update the sequence value to a new value. It
	// takes the expected current value of the sequence as well as the
	// new value to set it to. (true, nil) is returned if the value
	// was updated successfully. (false, nil) is returned if the
	// sequence was not at the expected value (indicating a concurrent
	// update). (false, <some error>) is returned for any other
	// problem.
	set(expected, next int) (bool, error)
}

const maxSeqRetries = 20

// updateSeqWithMin implements the abstract logic for incrementing a
// database backed sequence in a concurrency aware way.
//
// It is complicated because MongoDB's atomic update primitives don't
// provide a way to upsert a counter while also providing an initial
// value.  Instead, a number of database operations are used for each
// sequence update, relying on the atomicity guarantees that MongoDB
// offers. Optimistic database updates are attempted with retries when
// contention is observed.
func updateSeqWithMin(sequence seqUpdater, minVal int) (int, error) {
	for try := 0; try < maxSeqRetries; try++ {
		curVal, err := sequence.read()
		if err != nil {
			return -1, errors.Annotate(err, "could not read sequence")
		}
		if curVal == 0 {
			// No sequence document exists, create one.
			ok, err := sequence.create(minVal + 1)
			if err != nil {
				return -1, errors.Annotate(err, "could not create sequence")
			}
			if ok {
				return minVal, nil
			}
			// Someone else created the sequence document at the same
			// time, try again.
		} else {
			// Increment an existing sequence document, respecting the
			// minimum value provided.
			nextVal := curVal + 1
			if nextVal < minVal {
				nextVal = minVal + 1
			}
			ok, err := sequence.set(curVal, nextVal)
			if err != nil {
				return -1, errors.Annotate(err, "could not set sequence")
			}
			if ok {
				return nextVal - 1, nil
			}
			// Someone else incremented the sequence at the same time,
			// try again.
		}
	}
	return -1, errors.New("too much contention while updating sequence")
}

// dbSeqUpdater implements seqUpdater.
type dbSeqUpdater struct {
	coll mongo.Collection
	id   string
	name string
}

func (su *dbSeqUpdater) read() (int, error) {
	var doc bson.M
	err := su.coll.FindId(su.id).One(&doc)
	if errors.Cause(err) == mgo.ErrNotFound {
		return 0, nil
	} else if err != nil {
		return -1, errors.Trace(err)
	}
	return doc["counter"].(int), nil
}

func (su *dbSeqUpdater) create(value int) (bool, error) {
	err := su.coll.Writeable().Insert(bson.M{
		"_id":     su.id,
		"name":    su.name,
		"counter": value,
	})
	if mgo.IsDup(errors.Cause(err)) {
		return false, nil
	} else if err != nil {
		return false, errors.Trace(err)
	}
	return true, nil
}

func (su *dbSeqUpdater) set(expected, next int) (bool, error) {
	err := su.coll.Writeable().Update(
		bson.M{"_id": su.id, "counter": expected},
		bson.M{"$set": bson.M{"counter": next}},
	)
	if errors.Cause(err) == mgo.ErrNotFound {
		return false, nil
	} else if err != nil {
		return false, errors.Trace(err)
	}
	return true, nil
}<commit_msg>state: Added dbSeqUpdater.ensure<commit_after>// Copyright 2012, 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package state

import (
	"fmt"

	"github.com/juju/errors"

	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
)

type sequenceDoc struct {
	DocID     string `bson:"_id"`
	Name      string `bson:"name"`
	ModelUUID string `bson:"model-uuid"`
	Counter   int    // This stores the *next* value to return.
}

// sequence safely increments a database backed sequence, returning
// the next value.
func (st *State) sequence(name string) (int, error) {
	sequences, closer := st.getCollection(sequenceC)
	defer closer()
	query := sequences.FindId(name)
	inc := mgo.Change{
		Update: bson.M{
			"$set": bson.M{
				"name":       name,
				"model-uuid": st.ModelUUID(),
			},
			"$inc": bson.M{"counter": 1},
		},
		Upsert: true,
	}
	result := &sequenceDoc{}
	_, err := query.Apply(inc, result)
	if err != nil {
		return -1, fmt.Errorf("cannot increment %q sequence number: %v", name, err)
	}
	return result.Counter, nil
}

// sequenceWithMin safely increments a database backed sequence,
// allowing for a minimum value for the sequence to be specified. The
// minimum value is used as an initial value for the first use of a
// particular sequence. The minimum value will also cause a sequence
// value to jump ahead if the minimum is provided that is higher than
// the current sequence value.
//
// The data manipulated by `sequence` and `sequenceWithMin` is the
// same. It is safe to mix the 2 methods for the same sequence.
//
// `sequence` is more efficient than `sequenceWithMin` and should be
// preferred if there is no minimum value requirement.
func (st *State) sequenceWithMin(name string, minVal int) (int, error) {
	sequences, closer := st.getRawCollection(sequenceC)
	defer closer()
	updater := newDbSeqUpdater(sequences, st.ModelUUID(), name)
	return updateSeqWithMin(updater, minVal)
}

// seqUpdater abstracts away the database operations required for
// updating a sequence.
type seqUpdater interface {
	// read returns the current value of the sequence. If the sequence
	// doesn't exist yet it returns 0.
	read() (int, error)

	// create attempts to create a new sequence with the initial value
	// provided. It returns (true, nil) on success, (false, nil) if
	// the sequence already existed and (false, <some error>) if any
	// other error occurred.
	create(value int) (bool, error)

	// set attempts to update the sequence value to a new value. It
	// takes the expected current value of the sequence as well as the
	// new value to set it to. (true, nil) is returned if the value
	// was updated successfully. (false, nil) is returned if the
	// sequence was not at the expected value (indicating a concurrent
	// update). (false, <some error>) is returned for any other
	// problem.
	set(expected, next int) (bool, error)
}

const maxSeqRetries = 20

// updateSeqWithMin implements the abstract logic for incrementing a
// database backed sequence in a concurrency aware way.
//
// It is complicated because MongoDB's atomic update primitives don't
// provide a way to upsert a counter while also providing an initial
// value.  Instead, a number of database operations are used for each
// sequence update, relying on the atomicity guarantees that MongoDB
// offers. Optimistic database updates are attempted with retries when
// contention is observed.
func updateSeqWithMin(sequence seqUpdater, minVal int) (int, error) {
	for try := 0; try < maxSeqRetries; try++ {
		curVal, err := sequence.read()
		if err != nil {
			return -1, errors.Annotate(err, "could not read sequence")
		}
		if curVal == 0 {
			// No sequence document exists, create one.
			ok, err := sequence.create(minVal + 1)
			if err != nil {
				return -1, errors.Annotate(err, "could not create sequence")
			}
			if ok {
				return minVal, nil
			}
			// Someone else created the sequence document at the same
			// time, try again.
		} else {
			// Increment an existing sequence document, respecting the
			// minimum value provided.
			nextVal := curVal + 1
			if nextVal < minVal {
				nextVal = minVal + 1
			}
			ok, err := sequence.set(curVal, nextVal)
			if err != nil {
				return -1, errors.Annotate(err, "could not set sequence")
			}
			if ok {
				return nextVal - 1, nil
			}
			// Someone else incremented the sequence at the same time,
			// try again.
		}
	}
	return -1, errors.New("too much contention while updating sequence")
}

// dbSeqUpdater implements seqUpdater.
type dbSeqUpdater struct {
	coll      *mgo.Collection
	modelUUID string
	name      string
	id        string
}

func newDbSeqUpdater(coll *mgo.Collection, modelUUID, name string) *dbSeqUpdater {
	return &dbSeqUpdater{
		coll:      coll,
		modelUUID: modelUUID,
		name:      name,
		id:        modelUUID + ":" + name,
	}
}

func (su *dbSeqUpdater) read() (int, error) {
	var doc bson.M
	err := su.coll.FindId(su.id).One(&doc)
	if errors.Cause(err) == mgo.ErrNotFound {
		return 0, nil
	} else if err != nil {
		return -1, errors.Trace(err)
	}
	return doc["counter"].(int), nil
}

func (su *dbSeqUpdater) create(value int) (bool, error) {
	err := su.coll.Insert(bson.M{
		"_id":        su.id,
		"name":       su.name,
		"model-uuid": su.modelUUID,
		"counter":    value,
	})
	if mgo.IsDup(errors.Cause(err)) {
		return false, nil
	} else if err != nil {
		return false, errors.Trace(err)
	}
	return true, nil
}

func (su *dbSeqUpdater) set(expected, next int) (bool, error) {
	err := su.coll.Update(
		bson.M{"_id": su.id, "counter": expected},
		bson.M{"$set": bson.M{"counter": next}},
	)
	if errors.Cause(err) == mgo.ErrNotFound {
		return false, nil
	} else if err != nil {
		return false, errors.Trace(err)
	}
	return true, nil
}

func (su *dbSeqUpdater) ensure(next int) error {
	curVal, err := su.read()
	if err != nil {
		return errors.Trace(err)
	}

	var ok bool
	if curVal == 0 {
		ok, err = su.create(next)
	} else {
		// Sequences should never go backwards.
		if next <= curVal {
			return nil
		}
		ok, err = su.set(curVal, next)
	}
	if !ok {
		return errors.New("unexpected contention")
	}
	return errors.Trace(err)
}
{'fraction_non_alphanumeric': 0.11639667438073198, 'fraction_numerical': 0.005056998371475101, 'mean_word_length': 3.31189948263119, 'pattern_counts': {'":': 26, '<': 12, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7441358', 'n_tokens_mistral': 3992, 'n_tokens_neox': 3734, 'n_words': 1406}