<commit_before>// Go implementation of http://www.hashids.org under MIT license
// Generates hashes from an array of integers, eg. for YouTube like hashes
// Setup: go get github.com/speps/go-hashids
// Original implementations by Ivan Akimov at https://github.com/ivanakimov
// Thanks to Rémy Oudompheng and Peter Hellberg for code review and fixes

package hashids

import (
	"errors"
	"math"
)

const (
	// Version is the version number of the library
	Version string = "1.0.0"

	// DefaultAlphabet is the default alphabet used by go-hashids
	DefaultAlphabet string = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ1234567890"

	minAlphabetLength int     = 16
	sepDiv            float64 = 3.5
	guardDiv          float64 = 12.0
)

var sepsOriginal = []rune("cfhistuCFHISTU")

// HashID contains everything needed to encode/decode hashids
type HashID struct {
	alphabet  []rune
	minLength int
	salt      []rune
	seps      []rune
	guards    []rune
}

// HashIDData contains the information needed to generate hashids
type HashIDData struct {
	// Alphabet is the alphabet used to generate new ids
	Alphabet string

	// MinLength is the minimum length of a generated id
	MinLength int

	// Salt is the secret used to make the generated id harder to guess
	Salt string
}

// NewData creates a new HashIDData with the DefaultAlphabet already set.
func NewData() *HashIDData {
	return &HashIDData{Alphabet: DefaultAlphabet}
}

// New creates a new HashID
func New() *HashID {
	return NewWithData(NewData())
}

// NewWithData creates a new HashID with the provided HashIDData
func NewWithData(data *HashIDData) *HashID {
	if len(data.Alphabet) < minAlphabetLength {
		panic(errors.New("alphabet must contain at least 16 characters"))
	}
	// Check if all characters are unique in Alphabet
	uniqueCheck := make(map[rune]bool, len(data.Alphabet))
	for _, a := range data.Alphabet {
		if _, found := uniqueCheck[a]; found {
			panic(errors.New("duplicate character in alphabet"))
		}
		uniqueCheck[a] = true
	}

	alphabet := []rune(data.Alphabet)
	salt := []rune(data.Salt)

	seps := make([]rune, len(sepsOriginal))
	copy(seps, sepsOriginal)

	// seps should contain only characters present in alphabet; alphabet should not contains seps
	for i := 0; i < len(seps); i++ {
		foundIndex := -1
		for j, a := range alphabet {
			if a == seps[i] {
				foundIndex = j
				break
			}
		}
		if foundIndex == -1 {
			seps = append(seps[:i], seps[i+1:]...)
			i--
		} else {
			alphabet = append(alphabet[:foundIndex], alphabet[foundIndex+1:]...)
		}
	}
	seps = consistentShuffle(seps, salt)

	if len(seps) == 0 || float64(len(alphabet))/float64(len(seps)) > sepDiv {
		sepsLength := int(math.Ceil(float64(len(alphabet)) / sepDiv))
		if sepsLength == 1 {
			sepsLength++
		}
		if sepsLength > len(seps) {
			diff := sepsLength - len(seps)
			seps = append(seps, alphabet[:diff]...)
			alphabet = alphabet[diff:]
		} else {
			seps = seps[:sepsLength]
		}
	}
	alphabet = consistentShuffle(alphabet, salt)

	guardCount := int(math.Ceil(float64(len(alphabet)) / guardDiv))
	var guards []rune
	if len(alphabet) < 3 {
		guards = seps[:guardCount]
		seps = seps[guardCount:]
	} else {
		guards = alphabet[:guardCount]
		alphabet = alphabet[guardCount:]
	}

	return &HashID{
		alphabet:  alphabet,
		minLength: data.MinLength,
		salt:      salt,
		seps:      seps,
		guards:    guards,
	}
}

// Encode hashes an array of int to a string containing at least MinLength characters taken from the Alphabet.
// Use Decode using the same Alphabet and Salt to get back the array of int.
func (h *HashID) Encode(numbers []int) (string, error) {
	numbers64 := make([]int64, 0, len(numbers))
	for _, id := range numbers {
		numbers64 = append(numbers64, int64(id))
	}
	return h.EncodeInt64(numbers64)
}

// EncodeInt64 hashes an array of int64 to a string containing at least MinLength characters taken from the Alphabet.
// Use DecodeInt64 using the same Alphabet and Salt to get back the array of int64.
func (h *HashID) EncodeInt64(numbers []int64) (string, error) {
	if len(numbers) == 0 {
		return "", errors.New("encoding empty array of numbers makes no sense")
	}
	for _, n := range numbers {
		if n < 0 {
			return "", errors.New("negative number not supported")
		}
	}

	alphabet := make([]rune, len(h.alphabet))
	copy(alphabet, h.alphabet)

	numbersHash := int64(0)
	for i, n := range numbers {
		numbersHash += (n % int64(i+100))
	}

	result := make([]rune, 0, h.minLength)
	lottery := alphabet[numbersHash%int64(len(alphabet))]
	result = append(result, lottery)

	for i, n := range numbers {
		buffer := append([]rune{lottery}, append(h.salt, alphabet...)...)
		alphabet = consistentShuffle(alphabet, buffer[:len(alphabet)])
		hash := hash(n, alphabet)
		result = append(result, hash...)

		if i+1 < len(numbers) {
			n %= int64(hash[0]) + int64(i)
			result = append(result, h.seps[n%int64(len(h.seps))])
		}
	}

	if len(result) < h.minLength {
		guardIndex := (numbersHash + int64(result[0])) % int64(len(h.guards))
		result = append([]rune{h.guards[guardIndex]}, result...)

		if len(result) < h.minLength {
			guardIndex = (numbersHash + int64(result[2])) % int64(len(h.guards))
			result = append(result, h.guards[guardIndex])
		}
	}

	halfLength := len(alphabet) / 2
	for len(result) < h.minLength {
		alphabet = consistentShuffle(alphabet, alphabet)
		result = append(alphabet[halfLength:], append(result, alphabet[:halfLength]...)...)
		excess := len(result) - h.minLength
		if excess > 0 {
			result = result[excess/2 : excess/2+h.minLength]
		}
	}

	return string(result), nil
}

// DEPRECATED: Use DecryptWithError instead
// Decode unhashes the string passed to an array of int.
// It is symmetric with Encode if the Alphabet and Salt are the same ones which were used to hash.
// MinLength has no effect on Decode.
func (h *HashID) Decode(hash string) []int {
	result, err := h.DecodeWithError(hash)
	if err != nil {
		panic(err)
	}
	return result
}

// Decode unhashes the string passed to an array of int.
// It is symmetric with Encode if the Alphabet and Salt are the same ones which were used to hash.
// MinLength has no effect on Decode.
func (h *HashID) DecodeWithError(hash string) ([]int, error) {
	result64, err := h.DecodeInt64WithError(hash)
	if err != nil {
		return nil, err
	}
	result := make([]int, 0, len(result64))
	for _, id := range result64 {
		result = append(result, int(id))
	}
	return result, nil
}

// DEPRECATED: Use DecryptInt64WithError instead
// DecodeInt64 unhashes the string passed to an array of int64.
// It is symmetric with EncodeInt64 if the Alphabet and Salt are the same ones which were used to hash.
// MinLength has no effect on DecodeInt64.
func (h *HashID) DecodeInt64(hash string) []int64 {
	result, err := h.DecodeInt64WithError(hash)
	if err != nil {
		panic(err)
	}
	return result
}

// DecodeInt64 unhashes the string passed to an array of int64.
// It is symmetric with EncodeInt64 if the Alphabet and Salt are the same ones which were used to hash.
// MinLength has no effect on DecodeInt64.
func (h *HashID) DecodeInt64WithError(hash string) ([]int64, error) {
	hashes := splitRunes([]rune(hash), h.guards)
	hashIndex := 0
	if len(hashes) == 2 || len(hashes) == 3 {
		hashIndex = 1
	}

	result := make([]int64, 0)

	hashBreakdown := hashes[hashIndex]
	if len(hashBreakdown) > 0 {
		lottery := hashBreakdown[0]
		hashBreakdown = hashBreakdown[1:]
		hashes = splitRunes(hashBreakdown, h.seps)
		alphabet := []rune(h.alphabet)
		for _, subHash := range hashes {
			buffer := append([]rune{lottery}, append(h.salt, alphabet...)...)
			alphabet = consistentShuffle(alphabet, buffer[:len(alphabet)])
			number, err := unhash(subHash, alphabet)
			if err != nil {
				return nil, err
			}
			result = append(result, number)
		}
	}

	return result, nil
}

func splitRunes(input, seps []rune) [][]rune {
	splitIndices := make([]int, 0)
	for i, inputRune := range input {
		for _, sepsRune := range seps {
			if inputRune == sepsRune {
				splitIndices = append(splitIndices, i)
			}
		}
	}

	result := make([][]rune, 0, len(splitIndices)+1)
	inputLeft := input[:]
	for _, splitIndex := range splitIndices {
		splitIndex -= len(input) - len(inputLeft)
		subInput := make([]rune, splitIndex)
		copy(subInput, inputLeft[:splitIndex])
		result = append(result, subInput)
		inputLeft = inputLeft[splitIndex+1:]
	}
	result = append(result, inputLeft)

	return result
}

func hash(input int64, alphabet []rune) []rune {
	result := make([]rune, 0)
	for {
		r := alphabet[input%int64(len(alphabet))]
		result = append(result, r)
		input /= int64(len(alphabet))
		if input == 0 {
			break
		}
	}
	reversed := make([]rune, len(result))
	for i, r := range result {
		reversed[len(result)-i-1] = r
	}
	return reversed
}

func unhash(input, alphabet []rune) (int64, error) {
	result := int64(0)
	for i, inputRune := range input {
		alphabetPos := -1
		for pos, alphabetRune := range alphabet {
			if inputRune == alphabetRune {
				alphabetPos = pos
				break
			}
		}
		if alphabetPos == -1 {
			return 0, errors.New("alphabet used for hash was different")
		}

		result += int64(alphabetPos) * int64(math.Pow(float64(len(alphabet)), float64(len(input)-i-1)))
	}
	return result, nil
}

func consistentShuffle(alphabet, salt []rune) []rune {
	if len(salt) == 0 {
		return alphabet
	}

	result := make([]rune, len(alphabet))
	copy(result, alphabet)
	for i, v, p := len(result)-1, 0, 0; i > 0; i-- {
		p += int(salt[v])
		j := (int(salt[v]) + v + p) % i
		result[i], result[j] = result[j], result[i]
		v = (v + 1) % len(salt)
	}

	return result
}
<commit_msg>Don't use hard-coded value in string<commit_after>// Go implementation of http://www.hashids.org under MIT license
// Generates hashes from an array of integers, eg. for YouTube like hashes
// Setup: go get github.com/speps/go-hashids
// Original implementations by Ivan Akimov at https://github.com/ivanakimov
// Thanks to Rémy Oudompheng and Peter Hellberg for code review and fixes

package hashids

import (
	"errors"
	"math"
)

const (
	// Version is the version number of the library
	Version string = "1.0.0"

	// DefaultAlphabet is the default alphabet used by go-hashids
	DefaultAlphabet string = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ1234567890"

	minAlphabetLength int     = 16
	sepDiv            float64 = 3.5
	guardDiv          float64 = 12.0
)

var sepsOriginal = []rune("cfhistuCFHISTU")

// HashID contains everything needed to encode/decode hashids
type HashID struct {
	alphabet  []rune
	minLength int
	salt      []rune
	seps      []rune
	guards    []rune
}

// HashIDData contains the information needed to generate hashids
type HashIDData struct {
	// Alphabet is the alphabet used to generate new ids
	Alphabet string

	// MinLength is the minimum length of a generated id
	MinLength int

	// Salt is the secret used to make the generated id harder to guess
	Salt string
}

// NewData creates a new HashIDData with the DefaultAlphabet already set.
func NewData() *HashIDData {
	return &HashIDData{Alphabet: DefaultAlphabet}
}

// New creates a new HashID
func New() *HashID {
	return NewWithData(NewData())
}

// NewWithData creates a new HashID with the provided HashIDData
func NewWithData(data *HashIDData) *HashID {
	if len(data.Alphabet) < minAlphabetLength {
		panic(errors.New(fmt.Sprintf("alphabet must contain at least %d characters", minAlphabetLength)))
	}
	// Check if all characters are unique in Alphabet
	uniqueCheck := make(map[rune]bool, len(data.Alphabet))
	for _, a := range data.Alphabet {
		if _, found := uniqueCheck[a]; found {
			panic(errors.New("duplicate character in alphabet"))
		}
		uniqueCheck[a] = true
	}

	alphabet := []rune(data.Alphabet)
	salt := []rune(data.Salt)

	seps := make([]rune, len(sepsOriginal))
	copy(seps, sepsOriginal)

	// seps should contain only characters present in alphabet; alphabet should not contains seps
	for i := 0; i < len(seps); i++ {
		foundIndex := -1
		for j, a := range alphabet {
			if a == seps[i] {
				foundIndex = j
				break
			}
		}
		if foundIndex == -1 {
			seps = append(seps[:i], seps[i+1:]...)
			i--
		} else {
			alphabet = append(alphabet[:foundIndex], alphabet[foundIndex+1:]...)
		}
	}
	seps = consistentShuffle(seps, salt)

	if len(seps) == 0 || float64(len(alphabet))/float64(len(seps)) > sepDiv {
		sepsLength := int(math.Ceil(float64(len(alphabet)) / sepDiv))
		if sepsLength == 1 {
			sepsLength++
		}
		if sepsLength > len(seps) {
			diff := sepsLength - len(seps)
			seps = append(seps, alphabet[:diff]...)
			alphabet = alphabet[diff:]
		} else {
			seps = seps[:sepsLength]
		}
	}
	alphabet = consistentShuffle(alphabet, salt)

	guardCount := int(math.Ceil(float64(len(alphabet)) / guardDiv))
	var guards []rune
	if len(alphabet) < 3 {
		guards = seps[:guardCount]
		seps = seps[guardCount:]
	} else {
		guards = alphabet[:guardCount]
		alphabet = alphabet[guardCount:]
	}

	return &HashID{
		alphabet:  alphabet,
		minLength: data.MinLength,
		salt:      salt,
		seps:      seps,
		guards:    guards,
	}
}

// Encode hashes an array of int to a string containing at least MinLength characters taken from the Alphabet.
// Use Decode using the same Alphabet and Salt to get back the array of int.
func (h *HashID) Encode(numbers []int) (string, error) {
	numbers64 := make([]int64, 0, len(numbers))
	for _, id := range numbers {
		numbers64 = append(numbers64, int64(id))
	}
	return h.EncodeInt64(numbers64)
}

// EncodeInt64 hashes an array of int64 to a string containing at least MinLength characters taken from the Alphabet.
// Use DecodeInt64 using the same Alphabet and Salt to get back the array of int64.
func (h *HashID) EncodeInt64(numbers []int64) (string, error) {
	if len(numbers) == 0 {
		return "", errors.New("encoding empty array of numbers makes no sense")
	}
	for _, n := range numbers {
		if n < 0 {
			return "", errors.New("negative number not supported")
		}
	}

	alphabet := make([]rune, len(h.alphabet))
	copy(alphabet, h.alphabet)

	numbersHash := int64(0)
	for i, n := range numbers {
		numbersHash += (n % int64(i+100))
	}

	result := make([]rune, 0, h.minLength)
	lottery := alphabet[numbersHash%int64(len(alphabet))]
	result = append(result, lottery)

	for i, n := range numbers {
		buffer := append([]rune{lottery}, append(h.salt, alphabet...)...)
		alphabet = consistentShuffle(alphabet, buffer[:len(alphabet)])
		hash := hash(n, alphabet)
		result = append(result, hash...)

		if i+1 < len(numbers) {
			n %= int64(hash[0]) + int64(i)
			result = append(result, h.seps[n%int64(len(h.seps))])
		}
	}

	if len(result) < h.minLength {
		guardIndex := (numbersHash + int64(result[0])) % int64(len(h.guards))
		result = append([]rune{h.guards[guardIndex]}, result...)

		if len(result) < h.minLength {
			guardIndex = (numbersHash + int64(result[2])) % int64(len(h.guards))
			result = append(result, h.guards[guardIndex])
		}
	}

	halfLength := len(alphabet) / 2
	for len(result) < h.minLength {
		alphabet = consistentShuffle(alphabet, alphabet)
		result = append(alphabet[halfLength:], append(result, alphabet[:halfLength]...)...)
		excess := len(result) - h.minLength
		if excess > 0 {
			result = result[excess/2 : excess/2+h.minLength]
		}
	}

	return string(result), nil
}

// DEPRECATED: Use DecryptWithError instead
// Decode unhashes the string passed to an array of int.
// It is symmetric with Encode if the Alphabet and Salt are the same ones which were used to hash.
// MinLength has no effect on Decode.
func (h *HashID) Decode(hash string) []int {
	result, err := h.DecodeWithError(hash)
	if err != nil {
		panic(err)
	}
	return result
}

// Decode unhashes the string passed to an array of int.
// It is symmetric with Encode if the Alphabet and Salt are the same ones which were used to hash.
// MinLength has no effect on Decode.
func (h *HashID) DecodeWithError(hash string) ([]int, error) {
	result64, err := h.DecodeInt64WithError(hash)
	if err != nil {
		return nil, err
	}
	result := make([]int, 0, len(result64))
	for _, id := range result64 {
		result = append(result, int(id))
	}
	return result, nil
}

// DEPRECATED: Use DecryptInt64WithError instead
// DecodeInt64 unhashes the string passed to an array of int64.
// It is symmetric with EncodeInt64 if the Alphabet and Salt are the same ones which were used to hash.
// MinLength has no effect on DecodeInt64.
func (h *HashID) DecodeInt64(hash string) []int64 {
	result, err := h.DecodeInt64WithError(hash)
	if err != nil {
		panic(err)
	}
	return result
}

// DecodeInt64 unhashes the string passed to an array of int64.
// It is symmetric with EncodeInt64 if the Alphabet and Salt are the same ones which were used to hash.
// MinLength has no effect on DecodeInt64.
func (h *HashID) DecodeInt64WithError(hash string) ([]int64, error) {
	hashes := splitRunes([]rune(hash), h.guards)
	hashIndex := 0
	if len(hashes) == 2 || len(hashes) == 3 {
		hashIndex = 1
	}

	result := make([]int64, 0)

	hashBreakdown := hashes[hashIndex]
	if len(hashBreakdown) > 0 {
		lottery := hashBreakdown[0]
		hashBreakdown = hashBreakdown[1:]
		hashes = splitRunes(hashBreakdown, h.seps)
		alphabet := []rune(h.alphabet)
		for _, subHash := range hashes {
			buffer := append([]rune{lottery}, append(h.salt, alphabet...)...)
			alphabet = consistentShuffle(alphabet, buffer[:len(alphabet)])
			number, err := unhash(subHash, alphabet)
			if err != nil {
				return nil, err
			}
			result = append(result, number)
		}
	}

	return result, nil
}

func splitRunes(input, seps []rune) [][]rune {
	splitIndices := make([]int, 0)
	for i, inputRune := range input {
		for _, sepsRune := range seps {
			if inputRune == sepsRune {
				splitIndices = append(splitIndices, i)
			}
		}
	}

	result := make([][]rune, 0, len(splitIndices)+1)
	inputLeft := input[:]
	for _, splitIndex := range splitIndices {
		splitIndex -= len(input) - len(inputLeft)
		subInput := make([]rune, splitIndex)
		copy(subInput, inputLeft[:splitIndex])
		result = append(result, subInput)
		inputLeft = inputLeft[splitIndex+1:]
	}
	result = append(result, inputLeft)

	return result
}

func hash(input int64, alphabet []rune) []rune {
	result := make([]rune, 0)
	for {
		r := alphabet[input%int64(len(alphabet))]
		result = append(result, r)
		input /= int64(len(alphabet))
		if input == 0 {
			break
		}
	}
	reversed := make([]rune, len(result))
	for i, r := range result {
		reversed[len(result)-i-1] = r
	}
	return reversed
}

func unhash(input, alphabet []rune) (int64, error) {
	result := int64(0)
	for i, inputRune := range input {
		alphabetPos := -1
		for pos, alphabetRune := range alphabet {
			if inputRune == alphabetRune {
				alphabetPos = pos
				break
			}
		}
		if alphabetPos == -1 {
			return 0, errors.New("alphabet used for hash was different")
		}

		result += int64(alphabetPos) * int64(math.Pow(float64(len(alphabet)), float64(len(input)-i-1)))
	}
	return result, nil
}

func consistentShuffle(alphabet, salt []rune) []rune {
	if len(salt) == 0 {
		return alphabet
	}

	result := make([]rune, len(alphabet))
	copy(result, alphabet)
	for i, v, p := len(result)-1, 0, 0; i > 0; i-- {
		p += int(salt[v])
		j := (int(salt[v]) + v + p) % i
		result[i], result[j] = result[j], result[i]
		v = (v + 1) % len(salt)
	}

	return result
}
<|endoftext|>


<commit_before>// Copyright 2016 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package sync

import (
	"sync/atomic"
	"unsafe"
)

// Map is like a Go map[interface{}]interface{} but is safe for concurrent use
// by multiple goroutines without additional locking or coordination.
// Loads, stores, and deletes run in amortized constant time.
//
// The Map type is specialized. Most code should use a plain Go map instead,
// with separate locking or coordination, for better type safety and to make it
// easier to maintain other invariants along with the map content.
//
// The Map type is optimized for two common use cases: (1) when the entry for a given
// key is only ever written once but read many times, as in caches that only grow,
// or (2) when multiple goroutines read, write, and overwrite entries for disjoint
// sets of keys. In these two cases, use of a Map may significantly reduce lock
// contention compared to a Go map paired with a separate Mutex or RWMutex.
//
// The zero Map is empty and ready for use. A Map must not be copied after first use.
//
// In the terminology of the Go memory model, Map arranges that a write operation
// “synchronizes before” any read operation that observes the effect of the write, where
// read and write operations are defined as follows.
// Load, LoadAndDelete, LoadOrStore are read operations;
// Delete, LoadAndDelete, and Store are write operations;
// and LoadOrStore is a write operation when it returns loaded set to false.
type Map struct {
	mu Mutex

	// read contains the portion of the map's contents that are safe for
	// concurrent access (with or without mu held).
	//
	// The read field itself is always safe to load, but must only be stored with
	// mu held.
	//
	// Entries stored in read may be updated concurrently without mu, but updating
	// a previously-expunged entry requires that the entry be copied to the dirty
	// map and unexpunged with mu held.
	read atomic.Value // readOnly

	// dirty contains the portion of the map's contents that require mu to be
	// held. To ensure that the dirty map can be promoted to the read map quickly,
	// it also includes all of the non-expunged entries in the read map.
	//
	// Expunged entries are not stored in the dirty map. An expunged entry in the
	// clean map must be unexpunged and added to the dirty map before a new value
	// can be stored to it.
	//
	// If the dirty map is nil, the next write to the map will initialize it by
	// making a shallow copy of the clean map, omitting stale entries.
	dirty map[any]*entry

	// misses counts the number of loads since the read map was last updated that
	// needed to lock mu to determine whether the key was present.
	//
	// Once enough misses have occurred to cover the cost of copying the dirty
	// map, the dirty map will be promoted to the read map (in the unamended
	// state) and the next store to the map will make a new dirty copy.
	misses int
}

// readOnly is an immutable struct stored atomically in the Map.read field.
type readOnly struct {
	m       map[any]*entry
	amended bool // true if the dirty map contains some key not in m.
}

// expunged is an arbitrary pointer that marks entries which have been deleted
// from the dirty map.
var expunged = unsafe.Pointer(new(any))

// An entry is a slot in the map corresponding to a particular key.
type entry struct {
	// p points to the interface{} value stored for the entry.
	//
	// If p == nil, the entry has been deleted, and either m.dirty == nil or
	// m.dirty[key] is e.
	//
	// If p == expunged, the entry has been deleted, m.dirty != nil, and the entry
	// is missing from m.dirty.
	//
	// Otherwise, the entry is valid and recorded in m.read.m[key] and, if m.dirty
	// != nil, in m.dirty[key].
	//
	// An entry can be deleted by atomic replacement with nil: when m.dirty is
	// next created, it will atomically replace nil with expunged and leave
	// m.dirty[key] unset.
	//
	// An entry's associated value can be updated by atomic replacement, provided
	// p != expunged. If p == expunged, an entry's associated value can be updated
	// only after first setting m.dirty[key] = e so that lookups using the dirty
	// map find the entry.
	p unsafe.Pointer // *interface{}
}

func newEntry(i any) *entry {
	return &entry{p: unsafe.Pointer(&i)}
}

// Load returns the value stored in the map for a key, or nil if no
// value is present.
// The ok result indicates whether value was found in the map.
func (m *Map) Load(key any) (value any, ok bool) {
	read, _ := m.read.Load().(readOnly)
	e, ok := read.m[key]
	if !ok && read.amended {
		m.mu.Lock()
		// Avoid reporting a spurious miss if m.dirty got promoted while we were
		// blocked on m.mu. (If further loads of the same key will not miss, it's
		// not worth copying the dirty map for this key.)
		read, _ = m.read.Load().(readOnly)
		e, ok = read.m[key]
		if !ok && read.amended {
			e, ok = m.dirty[key]
			// Regardless of whether the entry was present, record a miss: this key
			// will take the slow path until the dirty map is promoted to the read
			// map.
			m.missLocked()
		}
		m.mu.Unlock()
	}
	if !ok {
		return nil, false
	}
	return e.load()
}

func (e *entry) load() (value any, ok bool) {
	p := atomic.LoadPointer(&e.p)
	if p == nil || p == expunged {
		return nil, false
	}
	return *(*any)(p), true
}

// Store sets the value for a key.
func (m *Map) Store(key, value any) {
	read, _ := m.read.Load().(readOnly)
	if e, ok := read.m[key]; ok && e.tryStore(&value) {
		return
	}

	m.mu.Lock()
	read, _ = m.read.Load().(readOnly)
	if e, ok := read.m[key]; ok {
		if e.unexpungeLocked() {
			// The entry was previously expunged, which implies that there is a
			// non-nil dirty map and this entry is not in it.
			m.dirty[key] = e
		}
		e.storeLocked(&value)
	} else if e, ok := m.dirty[key]; ok {
		e.storeLocked(&value)
	} else {
		if !read.amended {
			// We're adding the first new key to the dirty map.
			// Make sure it is allocated and mark the read-only map as incomplete.
			m.dirtyLocked()
			m.read.Store(readOnly{m: read.m, amended: true})
		}
		m.dirty[key] = newEntry(value)
	}
	m.mu.Unlock()
}

// tryStore stores a value if the entry has not been expunged.
//
// If the entry is expunged, tryStore returns false and leaves the entry
// unchanged.
func (e *entry) tryStore(i *any) bool {
	for {
		p := atomic.LoadPointer(&e.p)
		if p == expunged {
			return false
		}
		if atomic.CompareAndSwapPointer(&e.p, p, unsafe.Pointer(i)) {
			return true
		}
	}
}

// unexpungeLocked ensures that the entry is not marked as expunged.
//
// If the entry was previously expunged, it must be added to the dirty map
// before m.mu is unlocked.
func (e *entry) unexpungeLocked() (wasExpunged bool) {
	return atomic.CompareAndSwapPointer(&e.p, expunged, nil)
}

// storeLocked unconditionally stores a value to the entry.
//
// The entry must be known not to be expunged.
func (e *entry) storeLocked(i *any) {
	atomic.StorePointer(&e.p, unsafe.Pointer(i))
}

// LoadOrStore returns the existing value for the key if present.
// Otherwise, it stores and returns the given value.
// The loaded result is true if the value was loaded, false if stored.
func (m *Map) LoadOrStore(key, value any) (actual any, loaded bool) {
	// Avoid locking if it's a clean hit.
	read, _ := m.read.Load().(readOnly)
	if e, ok := read.m[key]; ok {
		actual, loaded, ok := e.tryLoadOrStore(value)
		if ok {
			return actual, loaded
		}
	}

	m.mu.Lock()
	read, _ = m.read.Load().(readOnly)
	if e, ok := read.m[key]; ok {
		if e.unexpungeLocked() {
			m.dirty[key] = e
		}
		actual, loaded, _ = e.tryLoadOrStore(value)
	} else if e, ok := m.dirty[key]; ok {
		actual, loaded, _ = e.tryLoadOrStore(value)
		m.missLocked()
	} else {
		if !read.amended {
			// We're adding the first new key to the dirty map.
			// Make sure it is allocated and mark the read-only map as incomplete.
			m.dirtyLocked()
			m.read.Store(readOnly{m: read.m, amended: true})
		}
		m.dirty[key] = newEntry(value)
		actual, loaded = value, false
	}
	m.mu.Unlock()

	return actual, loaded
}

// tryLoadOrStore atomically loads or stores a value if the entry is not
// expunged.
//
// If the entry is expunged, tryLoadOrStore leaves the entry unchanged and
// returns with ok==false.
func (e *entry) tryLoadOrStore(i any) (actual any, loaded, ok bool) {
	p := atomic.LoadPointer(&e.p)
	if p == expunged {
		return nil, false, false
	}
	if p != nil {
		return *(*any)(p), true, true
	}

	// Copy the interface after the first load to make this method more amenable
	// to escape analysis: if we hit the "load" path or the entry is expunged, we
	// shouldn't bother heap-allocating.
	ic := i
	for {
		if atomic.CompareAndSwapPointer(&e.p, nil, unsafe.Pointer(&ic)) {
			return i, false, true
		}
		p = atomic.LoadPointer(&e.p)
		if p == expunged {
			return nil, false, false
		}
		if p != nil {
			return *(*any)(p), true, true
		}
	}
}

// LoadAndDelete deletes the value for a key, returning the previous value if any.
// The loaded result reports whether the key was present.
func (m *Map) LoadAndDelete(key any) (value any, loaded bool) {
	read, _ := m.read.Load().(readOnly)
	e, ok := read.m[key]
	if !ok && read.amended {
		m.mu.Lock()
		read, _ = m.read.Load().(readOnly)
		e, ok = read.m[key]
		if !ok && read.amended {
			e, ok = m.dirty[key]
			delete(m.dirty, key)
			// Regardless of whether the entry was present, record a miss: this key
			// will take the slow path until the dirty map is promoted to the read
			// map.
			m.missLocked()
		}
		m.mu.Unlock()
	}
	if ok {
		return e.delete()
	}
	return nil, false
}

// Delete deletes the value for a key.
func (m *Map) Delete(key any) {
	m.LoadAndDelete(key)
}

func (e *entry) delete() (value any, ok bool) {
	for {
		p := atomic.LoadPointer(&e.p)
		if p == nil || p == expunged {
			return nil, false
		}
		if atomic.CompareAndSwapPointer(&e.p, p, nil) {
			return *(*any)(p), true
		}
	}
}

// Range calls f sequentially for each key and value present in the map.
// If f returns false, range stops the iteration.
//
// Range does not necessarily correspond to any consistent snapshot of the Map's
// contents: no key will be visited more than once, but if the value for any key
// is stored or deleted concurrently (including by f), Range may reflect any
// mapping for that key from any point during the Range call. Range does not
// block other methods on the receiver; even f itself may call any method on m.
//
// Range may be O(N) with the number of elements in the map even if f returns
// false after a constant number of calls.
func (m *Map) Range(f func(key, value any) bool) {
	// We need to be able to iterate over all of the keys that were already
	// present at the start of the call to Range.
	// If read.amended is false, then read.m satisfies that property without
	// requiring us to hold m.mu for a long time.
	read, _ := m.read.Load().(readOnly)
	if read.amended {
		// m.dirty contains keys not in read.m. Fortunately, Range is already O(N)
		// (assuming the caller does not break out early), so a call to Range
		// amortizes an entire copy of the map: we can promote the dirty copy
		// immediately!
		m.mu.Lock()
		read, _ = m.read.Load().(readOnly)
		if read.amended {
			read = readOnly{m: m.dirty}
			m.read.Store(read)
			m.dirty = nil
			m.misses = 0
		}
		m.mu.Unlock()
	}

	for k, e := range read.m {
		v, ok := e.load()
		if !ok {
			continue
		}
		if !f(k, v) {
			break
		}
	}
}

func (m *Map) missLocked() {
	m.misses++
	if m.misses < len(m.dirty) {
		return
	}
	m.read.Store(readOnly{m: m.dirty})
	m.dirty = nil
	m.misses = 0
}

func (m *Map) dirtyLocked() {
	if m.dirty != nil {
		return
	}

	read, _ := m.read.Load().(readOnly)
	m.dirty = make(map[any]*entry, len(read.m))
	for k, e := range read.m {
		if !e.tryExpungeLocked() {
			m.dirty[k] = e
		}
	}
}

func (e *entry) tryExpungeLocked() (isExpunged bool) {
	p := atomic.LoadPointer(&e.p)
	for p == nil {
		if atomic.CompareAndSwapPointer(&e.p, nil, expunged) {
			return true
		}
		p = atomic.LoadPointer(&e.p)
	}
	return p == expunged
}
<commit_msg>sync: switch Map to use atomic.Pointer<commit_after>// Copyright 2016 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package sync

import (
	"sync/atomic"
	"unsafe"
)

// Map is like a Go map[interface{}]interface{} but is safe for concurrent use
// by multiple goroutines without additional locking or coordination.
// Loads, stores, and deletes run in amortized constant time.
//
// The Map type is specialized. Most code should use a plain Go map instead,
// with separate locking or coordination, for better type safety and to make it
// easier to maintain other invariants along with the map content.
//
// The Map type is optimized for two common use cases: (1) when the entry for a given
// key is only ever written once but read many times, as in caches that only grow,
// or (2) when multiple goroutines read, write, and overwrite entries for disjoint
// sets of keys. In these two cases, use of a Map may significantly reduce lock
// contention compared to a Go map paired with a separate Mutex or RWMutex.
//
// The zero Map is empty and ready for use. A Map must not be copied after first use.
//
// In the terminology of the Go memory model, Map arranges that a write operation
// “synchronizes before” any read operation that observes the effect of the write, where
// read and write operations are defined as follows.
// Load, LoadAndDelete, LoadOrStore are read operations;
// Delete, LoadAndDelete, and Store are write operations;
// and LoadOrStore is a write operation when it returns loaded set to false.
type Map struct {
	mu Mutex

	// read contains the portion of the map's contents that are safe for
	// concurrent access (with or without mu held).
	//
	// The read field itself is always safe to load, but must only be stored with
	// mu held.
	//
	// Entries stored in read may be updated concurrently without mu, but updating
	// a previously-expunged entry requires that the entry be copied to the dirty
	// map and unexpunged with mu held.
	read atomic.Pointer[readOnly]

	// dirty contains the portion of the map's contents that require mu to be
	// held. To ensure that the dirty map can be promoted to the read map quickly,
	// it also includes all of the non-expunged entries in the read map.
	//
	// Expunged entries are not stored in the dirty map. An expunged entry in the
	// clean map must be unexpunged and added to the dirty map before a new value
	// can be stored to it.
	//
	// If the dirty map is nil, the next write to the map will initialize it by
	// making a shallow copy of the clean map, omitting stale entries.
	dirty map[any]*entry

	// misses counts the number of loads since the read map was last updated that
	// needed to lock mu to determine whether the key was present.
	//
	// Once enough misses have occurred to cover the cost of copying the dirty
	// map, the dirty map will be promoted to the read map (in the unamended
	// state) and the next store to the map will make a new dirty copy.
	misses int
}

// readOnly is an immutable struct stored atomically in the Map.read field.
type readOnly struct {
	m       map[any]*entry
	amended bool // true if the dirty map contains some key not in m.
}

// expunged is an arbitrary pointer that marks entries which have been deleted
// from the dirty map.
var expunged = unsafe.Pointer(new(any))

// An entry is a slot in the map corresponding to a particular key.
type entry struct {
	// p points to the interface{} value stored for the entry.
	//
	// If p == nil, the entry has been deleted, and either m.dirty == nil or
	// m.dirty[key] is e.
	//
	// If p == expunged, the entry has been deleted, m.dirty != nil, and the entry
	// is missing from m.dirty.
	//
	// Otherwise, the entry is valid and recorded in m.read.m[key] and, if m.dirty
	// != nil, in m.dirty[key].
	//
	// An entry can be deleted by atomic replacement with nil: when m.dirty is
	// next created, it will atomically replace nil with expunged and leave
	// m.dirty[key] unset.
	//
	// An entry's associated value can be updated by atomic replacement, provided
	// p != expunged. If p == expunged, an entry's associated value can be updated
	// only after first setting m.dirty[key] = e so that lookups using the dirty
	// map find the entry.
	p unsafe.Pointer // *interface{}
}

func newEntry(i any) *entry {
	return &entry{p: unsafe.Pointer(&i)}
}

func (m *Map) loadReadOnly() readOnly {
	if p := m.read.Load(); p != nil {
		return *p
	}
	return readOnly{}
}

// Load returns the value stored in the map for a key, or nil if no
// value is present.
// The ok result indicates whether value was found in the map.
func (m *Map) Load(key any) (value any, ok bool) {
	read := m.loadReadOnly()
	e, ok := read.m[key]
	if !ok && read.amended {
		m.mu.Lock()
		// Avoid reporting a spurious miss if m.dirty got promoted while we were
		// blocked on m.mu. (If further loads of the same key will not miss, it's
		// not worth copying the dirty map for this key.)
		read = m.loadReadOnly()
		e, ok = read.m[key]
		if !ok && read.amended {
			e, ok = m.dirty[key]
			// Regardless of whether the entry was present, record a miss: this key
			// will take the slow path until the dirty map is promoted to the read
			// map.
			m.missLocked()
		}
		m.mu.Unlock()
	}
	if !ok {
		return nil, false
	}
	return e.load()
}

func (e *entry) load() (value any, ok bool) {
	p := atomic.LoadPointer(&e.p)
	if p == nil || p == expunged {
		return nil, false
	}
	return *(*any)(p), true
}

// Store sets the value for a key.
func (m *Map) Store(key, value any) {
	read := m.loadReadOnly()
	if e, ok := read.m[key]; ok && e.tryStore(&value) {
		return
	}

	m.mu.Lock()
	read = m.loadReadOnly()
	if e, ok := read.m[key]; ok {
		if e.unexpungeLocked() {
			// The entry was previously expunged, which implies that there is a
			// non-nil dirty map and this entry is not in it.
			m.dirty[key] = e
		}
		e.storeLocked(&value)
	} else if e, ok := m.dirty[key]; ok {
		e.storeLocked(&value)
	} else {
		if !read.amended {
			// We're adding the first new key to the dirty map.
			// Make sure it is allocated and mark the read-only map as incomplete.
			m.dirtyLocked()
			m.read.Store(&readOnly{m: read.m, amended: true})
		}
		m.dirty[key] = newEntry(value)
	}
	m.mu.Unlock()
}

// tryStore stores a value if the entry has not been expunged.
//
// If the entry is expunged, tryStore returns false and leaves the entry
// unchanged.
func (e *entry) tryStore(i *any) bool {
	for {
		p := atomic.LoadPointer(&e.p)
		if p == expunged {
			return false
		}
		if atomic.CompareAndSwapPointer(&e.p, p, unsafe.Pointer(i)) {
			return true
		}
	}
}

// unexpungeLocked ensures that the entry is not marked as expunged.
//
// If the entry was previously expunged, it must be added to the dirty map
// before m.mu is unlocked.
func (e *entry) unexpungeLocked() (wasExpunged bool) {
	return atomic.CompareAndSwapPointer(&e.p, expunged, nil)
}

// storeLocked unconditionally stores a value to the entry.
//
// The entry must be known not to be expunged.
func (e *entry) storeLocked(i *any) {
	atomic.StorePointer(&e.p, unsafe.Pointer(i))
}

// LoadOrStore returns the existing value for the key if present.
// Otherwise, it stores and returns the given value.
// The loaded result is true if the value was loaded, false if stored.
func (m *Map) LoadOrStore(key, value any) (actual any, loaded bool) {
	// Avoid locking if it's a clean hit.
	read := m.loadReadOnly()
	if e, ok := read.m[key]; ok {
		actual, loaded, ok := e.tryLoadOrStore(value)
		if ok {
			return actual, loaded
		}
	}

	m.mu.Lock()
	read = m.loadReadOnly()
	if e, ok := read.m[key]; ok {
		if e.unexpungeLocked() {
			m.dirty[key] = e
		}
		actual, loaded, _ = e.tryLoadOrStore(value)
	} else if e, ok := m.dirty[key]; ok {
		actual, loaded, _ = e.tryLoadOrStore(value)
		m.missLocked()
	} else {
		if !read.amended {
			// We're adding the first new key to the dirty map.
			// Make sure it is allocated and mark the read-only map as incomplete.
			m.dirtyLocked()
			m.read.Store(&readOnly{m: read.m, amended: true})
		}
		m.dirty[key] = newEntry(value)
		actual, loaded = value, false
	}
	m.mu.Unlock()

	return actual, loaded
}

// tryLoadOrStore atomically loads or stores a value if the entry is not
// expunged.
//
// If the entry is expunged, tryLoadOrStore leaves the entry unchanged and
// returns with ok==false.
func (e *entry) tryLoadOrStore(i any) (actual any, loaded, ok bool) {
	p := atomic.LoadPointer(&e.p)
	if p == expunged {
		return nil, false, false
	}
	if p != nil {
		return *(*any)(p), true, true
	}

	// Copy the interface after the first load to make this method more amenable
	// to escape analysis: if we hit the "load" path or the entry is expunged, we
	// shouldn't bother heap-allocating.
	ic := i
	for {
		if atomic.CompareAndSwapPointer(&e.p, nil, unsafe.Pointer(&ic)) {
			return i, false, true
		}
		p = atomic.LoadPointer(&e.p)
		if p == expunged {
			return nil, false, false
		}
		if p != nil {
			return *(*any)(p), true, true
		}
	}
}

// LoadAndDelete deletes the value for a key, returning the previous value if any.
// The loaded result reports whether the key was present.
func (m *Map) LoadAndDelete(key any) (value any, loaded bool) {
	read := m.loadReadOnly()
	e, ok := read.m[key]
	if !ok && read.amended {
		m.mu.Lock()
		read = m.loadReadOnly()
		e, ok = read.m[key]
		if !ok && read.amended {
			e, ok = m.dirty[key]
			delete(m.dirty, key)
			// Regardless of whether the entry was present, record a miss: this key
			// will take the slow path until the dirty map is promoted to the read
			// map.
			m.missLocked()
		}
		m.mu.Unlock()
	}
	if ok {
		return e.delete()
	}
	return nil, false
}

// Delete deletes the value for a key.
func (m *Map) Delete(key any) {
	m.LoadAndDelete(key)
}

func (e *entry) delete() (value any, ok bool) {
	for {
		p := atomic.LoadPointer(&e.p)
		if p == nil || p == expunged {
			return nil, false
		}
		if atomic.CompareAndSwapPointer(&e.p, p, nil) {
			return *(*any)(p), true
		}
	}
}

// Range calls f sequentially for each key and value present in the map.
// If f returns false, range stops the iteration.
//
// Range does not necessarily correspond to any consistent snapshot of the Map's
// contents: no key will be visited more than once, but if the value for any key
// is stored or deleted concurrently (including by f), Range may reflect any
// mapping for that key from any point during the Range call. Range does not
// block other methods on the receiver; even f itself may call any method on m.
//
// Range may be O(N) with the number of elements in the map even if f returns
// false after a constant number of calls.
func (m *Map) Range(f func(key, value any) bool) {
	// We need to be able to iterate over all of the keys that were already
	// present at the start of the call to Range.
	// If read.amended is false, then read.m satisfies that property without
	// requiring us to hold m.mu for a long time.
	read := m.loadReadOnly()
	if read.amended {
		// m.dirty contains keys not in read.m. Fortunately, Range is already O(N)
		// (assuming the caller does not break out early), so a call to Range
		// amortizes an entire copy of the map: we can promote the dirty copy
		// immediately!
		m.mu.Lock()
		read = m.loadReadOnly()
		if read.amended {
			read = readOnly{m: m.dirty}
			m.read.Store(&read)
			m.dirty = nil
			m.misses = 0
		}
		m.mu.Unlock()
	}

	for k, e := range read.m {
		v, ok := e.load()
		if !ok {
			continue
		}
		if !f(k, v) {
			break
		}
	}
}

func (m *Map) missLocked() {
	m.misses++
	if m.misses < len(m.dirty) {
		return
	}
	m.read.Store(&readOnly{m: m.dirty})
	m.dirty = nil
	m.misses = 0
}

func (m *Map) dirtyLocked() {
	if m.dirty != nil {
		return
	}

	read := m.loadReadOnly()
	m.dirty = make(map[any]*entry, len(read.m))
	for k, e := range read.m {
		if !e.tryExpungeLocked() {
			m.dirty[k] = e
		}
	}
}

func (e *entry) tryExpungeLocked() (isExpunged bool) {
	p := atomic.LoadPointer(&e.p)
	for p == nil {
		if atomic.CompareAndSwapPointer(&e.p, nil, expunged) {
			return true
		}
		p = atomic.LoadPointer(&e.p)
	}
	return p == expunged
}
<|endoftext|>


<commit_before>/*
    norx.go
    ------

    This file is part of the Go reference implementation of NORX.

    :copyright: (c) 2014 Philipp Jovanovic <philipp@jovanovic.io>
    :license: BSD (3-Clause), see LICENSE
*/

package aead

const (
    NORX_W = 64              // wordsize
    NORX_R = 4               // number of rounds
    NORX_D = 1               // parallelism degree
    NORX_A = NORX_W * 4      // tag size
    NORX_N = NORX_W * 2      // nonce size
    NORX_K = NORX_W * 4      // key size
    NORX_B = NORX_W * 16     // state size
    NORX_C = NORX_W * 6      // capacity size
    RATE = NORX_B - NORX_C   // rate size
    HEADER_TAG  = 1 << 0
    PAYLOAD_TAG = 1 << 1
    TRAILER_TAG = 1 << 2
    FINAL_TAG   = 1 << 3
    BRANCH_TAG  = 1 << 4
    MERGE_TAG   = 1 << 5
    R0, R1, R2, R3 = 8, 19, 40, 63 // rotation offsets
    u0, u1 = 0x243F6A8885A308D3, 0x13198A2E03707344
    u2, u3 = 0xA4093822299F31D0, 0x082EFA98EC4E6C89
    u4, u5 = 0xAE8858DC339325A1, 0x670A134EE52D7FA6
    u6, u7 = 0xC4316D80CD967541, 0xD21DFBF8B630B762
    u8, u9 = 0x375A18D261E7F892, 0x343D1F187D92285B
)

type state_t struct {
    s [16]uint64
}

func ROTR(x,c uint64) uint64 {
    return (x >> c | x << (NORX_W - c))
}

func H(x,y uint64) uint64 {
    return (x ^ y) ^ ((x & y) << 1)
}

func G(a,b,c,d uint64) (uint64,uint64,uint64,uint64) {

    a = H(a,b)
    d = ROTR(a ^ d, R0)
    c = H(c,d)
    b = ROTR(b ^ c, R1)
    a = H(a,b)
    d = ROTR(a ^ d, R2)
    c = H(c,d)
    b = ROTR(b ^ c, R3)
    return a,b,c,d
}

func F(s []uint64) {
    // Column step
    s[ 0], s[ 4], s[ 8], s[12] = G(s[ 0], s[ 4], s[ 8], s[12])
    s[ 1], s[ 5], s[ 9], s[13] = G(s[ 1], s[ 5], s[ 9], s[13])
    s[ 2], s[ 6], s[10], s[14] = G(s[ 2], s[ 6], s[10], s[14])
    s[ 3], s[ 7], s[11], s[15] = G(s[ 3], s[ 7], s[11], s[15])
    // Diagonal step
    s[ 0], s[ 5], s[10], s[15] = G(s[ 0], s[ 5], s[10], s[15])
    s[ 1], s[ 6], s[11], s[12] = G(s[ 1], s[ 6], s[11], s[12])
    s[ 2], s[ 7], s[ 8], s[13] = G(s[ 2], s[ 7], s[ 8], s[13])
    s[ 3], s[ 4], s[ 9], s[14] = G(s[ 3], s[ 4], s[ 9], s[14])
}

func permute(state *state_t) {
    s := state.s[:]
    for i:=0; i<NORX_R; i++ { F(s) }
}

func setup(state *state_t, k []uint8, n []uint8) {

    s := state.s[:]

    s[ 0] = u0
    s[ 1] = LOAD64(n[ 0: 8])
    s[ 2] = LOAD64(n[ 8:16])
    s[ 3] = u1

    s[ 4] = LOAD64(k[ 0: 8])
    s[ 5] = LOAD64(k[ 8:16])
    s[ 6] = LOAD64(k[16:24])
    s[ 7] = LOAD64(k[24:32])

    s[ 8] = u2
    s[ 9] = u3
    s[10] = u4
    s[11] = u5

    s[12] = u6
    s[13] = u7
    s[14] = u8
    s[15] = u9

    s[14] ^= (NORX_R << 26) | (NORX_D << 18) | (NORX_W << 10) | NORX_A
    permute(state)
}

func process_header(state *state_t, in []uint8, inlen uint64) {

    if inlen > 0 {
        i := uint64(0)
        n := BYTES64(RATE)
        for inlen >= n {
            absorb_block(state, in[n*i:n*(i+1)], HEADER_TAG)
            inlen -= n
            i++
        }
        absorb_lastblock(state, in[n*i:n*i+inlen], inlen, HEADER_TAG)
    }
}

func encrypt_msg(state *state_t, out []uint8, in []uint8, inlen uint64) {

    if inlen > 0 {
        n := BYTES64(RATE)
        i := uint64(0)
        for inlen >= n {
            encrypt_block(state, out[n*i:n*(i+1)], in[n*i:n*(i+1)])
            inlen -= n
            i++
        }
        encrypt_lastblock(state, out[n*i:n*i+inlen], in[n*i:n*i+inlen], inlen)
    }
}

func decrypt_msg(state *state_t, out []uint8, in []uint8, inlen uint64) {

    if inlen > 0 {
        n := BYTES64(RATE)
        i := uint64(0)
        for inlen >= n {
            decrypt_block(state, out[n*i:n*(i+1)], in[n*i:n*(i+1)])
            inlen -= n
            i++
        }
        decrypt_lastblock(state, out[n*i:n*i+inlen], in[n*i:n*i+inlen], inlen)
    }
}

func process_trailer(state *state_t, in []uint8, inlen uint64) {

    if inlen > 0 {
        i := uint64(0)
        n := BYTES64(RATE)
        for inlen >= n {
            absorb_block(state, in[n*i:n*(i+1)], TRAILER_TAG)
            inlen -= n
            i++
        }
        absorb_lastblock(state, in[n*i:n*i+inlen], inlen, TRAILER_TAG)
    }
}

func output_tag(state *state_t, tag []uint8) {

    inject_tag(state, FINAL_TAG)
    permute(state)
    permute(state)

    s := state.s[:]
    lastblock := make([]uint8, BYTES64(RATE))
    b := BYTES64(NORX_W)
    for i:= uint64(0); i < WORDS64(RATE); i++ {
        STORE64(lastblock[b*i:b*(i+1)], s[i])
    }
    copy(tag[:], lastblock[:])
    BURN8(lastblock[:], BYTES64(RATE))
}

func verify_tag(tag1 []uint8, tag2 []uint8) int {

    acc := 0
    for i := uint64(0); i < BYTES64(NORX_A); i++ {
        acc |= int(tag1[i] ^ tag2[i])
    }
    return (((acc - 1) >> 8) & 1) - 1
}

func pad(out []uint8, in []uint8, inlen uint64) {

    copy(out[:],in[:inlen])
    out[inlen] = 0x01
    out[BYTES64(RATE) - 1] |= 0x80
}

func inject_tag(state *state_t, tag uint64) {
    s := state.s[:]
    s[15] ^= tag
}

func absorb_block(state *state_t, in []uint8, tag uint64) {

    inject_tag(state, tag)
    permute(state)

    s := state.s[:]
    b := BYTES64(NORX_W)
    for i:= uint64(0); i < WORDS64(RATE); i++ {
        s[i] ^= LOAD64(in[b*i:b*(i+1)])
    }
}

func absorb_lastblock(state *state_t, in []uint8, inlen uint64, tag uint64) {

    lastblock := make([]uint8, BYTES64(RATE))
    pad(lastblock[:], in[:], inlen)
    absorb_block(state, lastblock[:], tag)
    BURN8(lastblock[:], BYTES64(RATE))
}

func encrypt_block(state *state_t, out []uint8, in []uint8) {

    inject_tag(state, PAYLOAD_TAG)
    permute(state)

    s := state.s[:]
    b := BYTES64(NORX_W)
    for i:= uint64(0); i < WORDS64(RATE); i++ {
        s[i] ^= LOAD64(in[b*i:b*(i+1)])
        STORE64(out[b*i:b*(i+1)], s[i])
    }
}

func encrypt_lastblock(state *state_t, out []uint8, in []uint8, inlen uint64) {

    lastblock := make([]uint8, BYTES64(RATE))
    pad(lastblock[:], in[:], inlen)
    encrypt_block(state, lastblock[:], lastblock[:])
    copy(out[:], lastblock[:])
    BURN8(lastblock[:], BYTES64(RATE))
}

func decrypt_block(state *state_t, out []uint8, in []uint8) {

    inject_tag(state, PAYLOAD_TAG)
    permute(state)

    s := state.s[:]
    b := BYTES64(NORX_W)
    for i:= uint64(0); i < WORDS64(RATE); i++ {
        c := LOAD64(in[b*i:b*(i+1)])
        STORE64(out[b*i:b*(i+1)], s[i] ^ c)
        s[i] = c
    }
}

func decrypt_lastblock(state *state_t, out []uint8, in []uint8, inlen uint64) {

    inject_tag(state, PAYLOAD_TAG)
    permute(state)

    i := uint64(0)
    s := state.s[:]
    n := BYTES64(NORX_W)
    b := make([]uint8, n)

    /* undo padding */
    s[inlen / n] ^= uint64(0x01) << (inlen % n * 8)
    s[WORDS64(RATE) - 1] ^= uint64(0x80) << ((BYTES64(RATE) - 1) % n * 8)

    for inlen >= n {
        c := LOAD64(in[n*i:n*(i+1)])
        STORE64(out[n*i:n*(i+1)], s[i] ^ c)
        s[i] = c
        inlen -= n
        i++
    }

    STORE64(b,s[i])
    for j := uint64(0); j < inlen; j++ {
      c := in[n*i+j]
      out[n*i+j] = b[j] ^ c
      b[j] = c
    }
    s[i] = LOAD64(b)
    BURN8(b,n)
}

func AEAD_encrypt(
    c []uint8, clen *uint64,
    h []uint8, hlen uint64,
    m []uint8, mlen uint64,
    t []uint8, tlen uint64,
    nonce []uint8,
    key []uint8) {

    state := new(state_t)
    setup(state, key, nonce)
    process_header(state, h, hlen)
    encrypt_msg(state, c, m, mlen)
    process_trailer(state, t, tlen)
    output_tag(state, c[ mlen: ])
    *clen = mlen + BYTES64(NORX_A)
    BURN64(state.s[:], WORDS64(NORX_B))
}

func AEAD_decrypt(
    m []uint8, mlen *uint64,
    h []uint8, hlen uint64,
    c []uint8, clen uint64,
    t []uint8, tlen uint64,
    nonce []uint8,
    key []uint8) int {

    if clen < BYTES64(NORX_A) {
      return -1
    }
    result := -1
    tag := make([]uint8, NORX_A)
    state := new(state_t)
    setup(state, key, nonce)
    process_header(state, h, hlen)
    decrypt_msg(state, m, c, clen - BYTES64(NORX_A))
    process_trailer(state, t, tlen)
    output_tag(state, tag)
    *mlen = clen - BYTES64(NORX_A)
    result = verify_tag(c[clen - BYTES64(NORX_A):], tag)
    if result != 0 {
        BURN8(m, clen - BYTES64(NORX_A))
    }
    BURN64(state.s[:], WORDS64(NORX_B))
    return result
}

<commit_msg>cleanup<commit_after>/*
    norx.go
    ------

    This file is part of the Go reference implementation of NORX.

    :copyright: (c) 2014 Philipp Jovanovic <philipp@jovanovic.io>
    :license: BSD (3-Clause), see LICENSE
*/

package aead

const (
    NORX_W = 64              // wordsize
    NORX_R = 4               // number of rounds
    NORX_D = 1               // parallelism degree
    NORX_A = NORX_W * 4      // tag size
    NORX_N = NORX_W * 2      // nonce size
    NORX_K = NORX_W * 4      // key size
    NORX_B = NORX_W * 16     // state size
    NORX_C = NORX_W * 6      // capacity size
    RATE = NORX_B - NORX_C   // rate size
    HEADER_TAG  = 1 << 0
    PAYLOAD_TAG = 1 << 1
    TRAILER_TAG = 1 << 2
    FINAL_TAG   = 1 << 3
    BRANCH_TAG  = 1 << 4
    MERGE_TAG   = 1 << 5
    R0, R1, R2, R3 = 8, 19, 40, 63 // rotation offsets
    u0, u1 = 0x243F6A8885A308D3, 0x13198A2E03707344
    u2, u3 = 0xA4093822299F31D0, 0x082EFA98EC4E6C89
    u4, u5 = 0xAE8858DC339325A1, 0x670A134EE52D7FA6
    u6, u7 = 0xC4316D80CD967541, 0xD21DFBF8B630B762
    u8, u9 = 0x375A18D261E7F892, 0x343D1F187D92285B
)

type state_t struct {
    s [16]uint64
}

func ROTR(x,c uint64) uint64 {
    return (x >> c | x << (NORX_W - c))
}

func H(x,y uint64) uint64 {
    return (x ^ y) ^ ((x & y) << 1)
}

func G(a,b,c,d uint64) (uint64,uint64,uint64,uint64) {

    a = H(a,b)
    d = ROTR(a ^ d, R0)
    c = H(c,d)
    b = ROTR(b ^ c, R1)
    a = H(a,b)
    d = ROTR(a ^ d, R2)
    c = H(c,d)
    b = ROTR(b ^ c, R3)
    return a,b,c,d
}

func F(s []uint64) {
    // Column step
    s[ 0], s[ 4], s[ 8], s[12] = G(s[ 0], s[ 4], s[ 8], s[12])
    s[ 1], s[ 5], s[ 9], s[13] = G(s[ 1], s[ 5], s[ 9], s[13])
    s[ 2], s[ 6], s[10], s[14] = G(s[ 2], s[ 6], s[10], s[14])
    s[ 3], s[ 7], s[11], s[15] = G(s[ 3], s[ 7], s[11], s[15])
    // Diagonal step
    s[ 0], s[ 5], s[10], s[15] = G(s[ 0], s[ 5], s[10], s[15])
    s[ 1], s[ 6], s[11], s[12] = G(s[ 1], s[ 6], s[11], s[12])
    s[ 2], s[ 7], s[ 8], s[13] = G(s[ 2], s[ 7], s[ 8], s[13])
    s[ 3], s[ 4], s[ 9], s[14] = G(s[ 3], s[ 4], s[ 9], s[14])
}

func permute(state *state_t) {
    s := state.s[:]
    for i:=0; i<NORX_R; i++ { F(s) }
}

func setup(state *state_t, k []uint8, n []uint8) {

    s := state.s[:]

    s[ 0] = u0
    s[ 1] = LOAD64(n[ 0: 8])
    s[ 2] = LOAD64(n[ 8:16])
    s[ 3] = u1

    s[ 4] = LOAD64(k[ 0: 8])
    s[ 5] = LOAD64(k[ 8:16])
    s[ 6] = LOAD64(k[16:24])
    s[ 7] = LOAD64(k[24:32])

    s[ 8] = u2
    s[ 9] = u3
    s[10] = u4
    s[11] = u5

    s[12] = u6
    s[13] = u7
    s[14] = u8
    s[15] = u9

    s[14] ^= (NORX_R << 26) | (NORX_D << 18) | (NORX_W << 10) | NORX_A
    permute(state)
}

func process_header(state *state_t, in []uint8, inlen uint64) {

    if inlen > 0 {
        i := uint64(0)
        n := BYTES64(RATE)
        for inlen >= n {
            absorb_block(state, in[n*i:n*(i+1)], HEADER_TAG)
            inlen -= n
            i++
        }
        absorb_lastblock(state, in[n*i:n*i+inlen], inlen, HEADER_TAG)
    }
}

func encrypt_msg(state *state_t, out []uint8, in []uint8, inlen uint64) {

    if inlen > 0 {
        i := uint64(0)
        n := BYTES64(RATE)
        for inlen >= n {
            encrypt_block(state, out[n*i:n*(i+1)], in[n*i:n*(i+1)])
            inlen -= n
            i++
        }
        encrypt_lastblock(state, out[n*i:n*i+inlen], in[n*i:n*i+inlen], inlen)
    }
}

func decrypt_msg(state *state_t, out []uint8, in []uint8, inlen uint64) {

    if inlen > 0 {
        i := uint64(0)
        n := BYTES64(RATE)
        for inlen >= n {
            decrypt_block(state, out[n*i:n*(i+1)], in[n*i:n*(i+1)])
            inlen -= n
            i++
        }
        decrypt_lastblock(state, out[n*i:n*i+inlen], in[n*i:n*i+inlen], inlen)
    }
}

func process_trailer(state *state_t, in []uint8, inlen uint64) {

    if inlen > 0 {
        i := uint64(0)
        n := BYTES64(RATE)
        for inlen >= n {
            absorb_block(state, in[n*i:n*(i+1)], TRAILER_TAG)
            inlen -= n
            i++
        }
        absorb_lastblock(state, in[n*i:n*i+inlen], inlen, TRAILER_TAG)
    }
}

func output_tag(state *state_t, tag []uint8) {

    inject_tag(state, FINAL_TAG)
    permute(state)
    permute(state)

    s := state.s[:]
    lastblock := make([]uint8, BYTES64(RATE))
    b := BYTES64(NORX_W)
    for i:= uint64(0); i < WORDS64(RATE); i++ {
        STORE64(lastblock[b*i:b*(i+1)], s[i])
    }
    copy(tag[:], lastblock[:])
    BURN8(lastblock[:], BYTES64(RATE))
}

func verify_tag(tag1 []uint8, tag2 []uint8) int {

    acc := 0
    for i := uint64(0); i < BYTES64(NORX_A); i++ {
        acc |= int(tag1[i] ^ tag2[i])
    }
    return (((acc - 1) >> 8) & 1) - 1
}

func pad(out []uint8, in []uint8, inlen uint64) {

    copy(out[:],in[:inlen])
    out[inlen] = 0x01
    out[BYTES64(RATE) - 1] |= 0x80
}

func inject_tag(state *state_t, tag uint64) {
    s := state.s[:]
    s[15] ^= tag
}

func absorb_block(state *state_t, in []uint8, tag uint64) {

    inject_tag(state, tag)
    permute(state)

    s := state.s[:]
    b := BYTES64(NORX_W)
    for i:= uint64(0); i < WORDS64(RATE); i++ {
        s[i] ^= LOAD64(in[b*i:b*(i+1)])
    }
}

func absorb_lastblock(state *state_t, in []uint8, inlen uint64, tag uint64) {

    lastblock := make([]uint8, BYTES64(RATE))
    pad(lastblock[:], in[:], inlen)
    absorb_block(state, lastblock[:], tag)
    BURN8(lastblock[:], BYTES64(RATE))
}

func encrypt_block(state *state_t, out []uint8, in []uint8) {

    inject_tag(state, PAYLOAD_TAG)
    permute(state)

    s := state.s[:]
    b := BYTES64(NORX_W)
    for i:= uint64(0); i < WORDS64(RATE); i++ {
        s[i] ^= LOAD64(in[b*i:b*(i+1)])
        STORE64(out[b*i:b*(i+1)], s[i])
    }
}

func encrypt_lastblock(state *state_t, out []uint8, in []uint8, inlen uint64) {

    lastblock := make([]uint8, BYTES64(RATE))
    pad(lastblock[:], in[:], inlen)
    encrypt_block(state, lastblock[:], lastblock[:])
    copy(out[:], lastblock[:])
    BURN8(lastblock[:], BYTES64(RATE))
}

func decrypt_block(state *state_t, out []uint8, in []uint8) {

    inject_tag(state, PAYLOAD_TAG)
    permute(state)

    s := state.s[:]
    b := BYTES64(NORX_W)
    for i:= uint64(0); i < WORDS64(RATE); i++ {
        c := LOAD64(in[b*i:b*(i+1)])
        STORE64(out[b*i:b*(i+1)], s[i] ^ c)
        s[i] = c
    }
}

func decrypt_lastblock(state *state_t, out []uint8, in []uint8, inlen uint64) {

    inject_tag(state, PAYLOAD_TAG)
    permute(state)

    s := state.s[:]
    i := uint64(0)
    n := BYTES64(NORX_W)
    b := make([]uint8, n)

    /* undo padding */
    s[inlen / n] ^= uint64(0x01) << (inlen % n * 8)
    s[WORDS64(RATE) - 1] ^= uint64(0x80) << ((BYTES64(RATE) - 1) % n * 8)

    for inlen >= n {
        c := LOAD64(in[n*i:n*(i+1)])
        STORE64(out[n*i:n*(i+1)], s[i] ^ c)
        s[i] = c
        inlen -= n
        i++
    }

    STORE64(b,s[i])
    for j := uint64(0); j < inlen; j++ {
      c := in[n*i+j]
      out[n*i+j] = b[j] ^ c
      b[j] = c
    }
    s[i] = LOAD64(b)
    BURN8(b,n)
}

func AEAD_encrypt(
    c []uint8, clen *uint64,
    h []uint8, hlen uint64,
    m []uint8, mlen uint64,
    t []uint8, tlen uint64,
    nonce []uint8,
    key []uint8) {

    state := new(state_t)
    setup(state, key, nonce)
    process_header(state, h, hlen)
    encrypt_msg(state, c, m, mlen)
    process_trailer(state, t, tlen)
    output_tag(state, c[ mlen: ])
    *clen = mlen + BYTES64(NORX_A)
    BURN64(state.s[:], WORDS64(NORX_B))
}

func AEAD_decrypt(
    m []uint8, mlen *uint64,
    h []uint8, hlen uint64,
    c []uint8, clen uint64,
    t []uint8, tlen uint64,
    nonce []uint8,
    key []uint8) int {

    if clen < BYTES64(NORX_A) {
      return -1
    }
    result := -1
    tag := make([]uint8, NORX_A)
    state := new(state_t)
    setup(state, key, nonce)
    process_header(state, h, hlen)
    decrypt_msg(state, m, c, clen - BYTES64(NORX_A))
    process_trailer(state, t, tlen)
    output_tag(state, tag)
    *mlen = clen - BYTES64(NORX_A)
    result = verify_tag(c[clen - BYTES64(NORX_A):], tag)
    if result != 0 {
        BURN8(m, clen - BYTES64(NORX_A))
    }
    BURN64(state.s[:], WORDS64(NORX_B))
    return result
}

<|endoftext|>


<commit_before>/*******************************************************************************
 * Copyright 2014 by Artem Andreenko, Openstat.
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to
 * deal in the Software without restriction, including without limitation the
 * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
 * sell copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
 * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
 * IN THE SOFTWARE.
 ******************************************************************************/

package goaero

// #cgo CFLAGS: -I/usr/local/include
// #cgo LDFLAGS: /usr/local/lib/libaerospike.a -llua -lcrypto -lz
// #include <stdlib.h>
// #include <aerospike/aerospike.h>
// #include <aerospike/aerospike_key.h>
// #include <aerospike/aerospike_scan.h>
// #include <aerospike/as_log.h>
import "C"

import "unsafe"

type Aerospike struct {
	aerospike C.aerospike
	config    *Config
}

func NewAerospike(config *Config) (self *Aerospike) {
	self = new(Aerospike)
	self.config = config
	C.aerospike_init(&self.aerospike, self.config.getCStruct())
	return
}

func (self *Aerospike) Connect() (err error) {
	var e C.as_error
	if C.aerospike_connect(&self.aerospike, &e) != C.AEROSPIKE_OK {
		return as_error(e)
	}
	return
}

func (self *Aerospike) Close() (err error) {
	var e C.as_error
	if C.aerospike_close(&self.aerospike, &e) != C.AEROSPIKE_OK {
		return as_error(e)
	}
	C.aerospike_destroy(&self.aerospike)
	return
}

func (self *Aerospike) Put(key *Key, rec *Record, policy_write *PolicyWrite) (err error) {
	var e C.as_error
	var policy *C.as_policy_write
	if policy_write != nil {
		policy = &policy_write.as_policy_write
	}
	if C.aerospike_key_put(&self.aerospike, &e, policy, &key.as_key, rec.p_as_record) != C.AEROSPIKE_OK {
		return as_error(e)
	}
	return
}

func (self *Aerospike) Get(key *Key, rec *Record, policy_read *PolicyRead) (err error) {
	var e C.as_error
	var policy *C.as_policy_read
	if policy_read != nil {
		policy = &policy_read.as_policy_read
	}
	if C.aerospike_key_get(&self.aerospike, &e, policy, &key.as_key, &rec.p_as_record) != C.AEROSPIKE_OK {
		return as_error(e)
	}
	return
}

func (self * Aerospike) KeyOperate(key * Key, ops * Operations, rec * Record, policy_operate * PolicyOperate) (err error) {
	var e C.as_error
	var policy *C.as_policy_operate
	if policy_operate != nil {
		policy = &policy_operate.as_policy_operate
	}
	if C.aerospike_key_operate(&self.aerospike, &e, policy, &key.as_key, &ops.as_ops, &rec.p_as_record) != C.AEROSPIKE_OK {
		return as_error(e)
	}
	return
}

// callback: func (* C.as_val, unsafe.Pointer) bool
func (self * Aerospike) ScanForeach(scan * Scan, callback * [0]byte, udata unsafe.Pointer, policy_scan * PolicyScan) (err error) {
	var e C.as_error
	var policy * C.as_policy_scan
	if policy_scan != nil {
		policy = &policy_scan.as_policy_scan
	}
	if C.aerospike_scan_foreach(&self.aerospike, &e, policy, &scan.as_scan, callback, udata) != C.AEROSPIKE_OK {
		return as_error(e)
	}
	return
}
<commit_msg>Close() double free bugfix<commit_after>/*******************************************************************************
 * Copyright 2014 by Artem Andreenko, Openstat.
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to
 * deal in the Software without restriction, including without limitation the
 * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
 * sell copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
 * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
 * IN THE SOFTWARE.
 ******************************************************************************/

package goaero

// #cgo CFLAGS: -I/usr/local/include
// #cgo LDFLAGS: /usr/local/lib/libaerospike.a -llua -lcrypto -lz
// #include <stdlib.h>
// #include <aerospike/aerospike.h>
// #include <aerospike/aerospike_key.h>
// #include <aerospike/aerospike_scan.h>
// #include <aerospike/as_log.h>
import "C"

import "unsafe"

type Aerospike struct {
	aerospike C.aerospike
	config    *Config
}

func NewAerospike(config *Config) (self *Aerospike) {
	self = new(Aerospike)
	self.config = config
	C.aerospike_init(&self.aerospike, self.config.getCStruct())
	return
}

func (self *Aerospike) Connect() (err error) {
	var e C.as_error
	if C.aerospike_connect(&self.aerospike, &e) != C.AEROSPIKE_OK {
		return as_error(e)
	}
	return
}

func (self *Aerospike) Close() (err error) {
	var e C.as_error
	if C.aerospike_close(&self.aerospike, &e) != C.AEROSPIKE_OK {
		return as_error(e)
	} else {
		C.aerospike_destroy(&self.aerospike)
	}
	return
}

func (self *Aerospike) Put(key *Key, rec *Record, policy_write *PolicyWrite) (err error) {
	var e C.as_error
	var policy *C.as_policy_write
	if policy_write != nil {
		policy = &policy_write.as_policy_write
	}
	if C.aerospike_key_put(&self.aerospike, &e, policy, &key.as_key, rec.p_as_record) != C.AEROSPIKE_OK {
		return as_error(e)
	}
	return
}

func (self *Aerospike) Get(key *Key, rec *Record, policy_read *PolicyRead) (err error) {
	var e C.as_error
	var policy *C.as_policy_read
	if policy_read != nil {
		policy = &policy_read.as_policy_read
	}
	if C.aerospike_key_get(&self.aerospike, &e, policy, &key.as_key, &rec.p_as_record) != C.AEROSPIKE_OK {
		return as_error(e)
	}
	return
}

func (self * Aerospike) KeyOperate(key * Key, ops * Operations, rec * Record, policy_operate * PolicyOperate) (err error) {
	var e C.as_error
	var policy *C.as_policy_operate
	if policy_operate != nil {
		policy = &policy_operate.as_policy_operate
	}
	if C.aerospike_key_operate(&self.aerospike, &e, policy, &key.as_key, &ops.as_ops, &rec.p_as_record) != C.AEROSPIKE_OK {
		return as_error(e)
	}
	return
}

// callback: func (* C.as_val, unsafe.Pointer) bool
func (self * Aerospike) ScanForeach(scan * Scan, callback * [0]byte, udata unsafe.Pointer, policy_scan * PolicyScan) (err error) {
	var e C.as_error
	var policy * C.as_policy_scan
	if policy_scan != nil {
		policy = &policy_scan.as_policy_scan
	}
	if C.aerospike_scan_foreach(&self.aerospike, &e, policy, &scan.as_scan, callback, udata) != C.AEROSPIKE_OK {
		return as_error(e)
	}
	return
}
<|endoftext|>


<commit_before>/*
 * Copyright (c) 2012 Matt Jibson <matt.jibson@gmail.com>
 *
 * Permission to use, copy, modify, and distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */

package goon

import (
	"appengine/datastore"
	"reflect"
)

// Count returns the number of results for the query.
func (g *Goon) Count(q *datastore.Query) (int, error) {
	return q.Count(g.context)
}

// GetAll returns the entities in the given query. Returned entities (but not
// the query itself) are cached in local memory (but not memcache).
//
// Returns nil if q is a keys-only query or dst is not a slice.
//
// Refer to appengine/datastore.Query.Count:
// https://developers.google.com/appengine/docs/go/datastore/reference#Query.GetAll
//
// Example usage:
//   g := goon.NewGoon(r)
//   q := datastore.NewQuery("Group")
//   var gg []Group // (or []*Group)
//   es, err := g.GetAll(q, &gg)
func (g *Goon) GetAll(q *datastore.Query, dst interface{}) ([]*Entity, error) {
	keys, err := q.GetAll(g.context, dst)
	if err != nil {
		g.error(err)
		return nil, err
	}

	keysOnly := dst == nil

	var v reflect.Value
	var t reflect.Type
	if !keysOnly {
		v = reflect.Indirect(reflect.ValueOf(dst))
		t = v.Type()

		// try to detect a keys-only query
		if t.Kind() != reflect.Slice || v.Len() != len(keys) {
			keysOnly = true
		}
	}

	es := make([]*Entity, len(keys))

	for i, k := range keys {
		var e *Entity
		if keysOnly {
			e = NewEntity(k, nil)
		} else {
			e = NewEntity(k, v.Index(i).Interface())
		}
		es[i] = e

		if !g.inTransaction {
			g.cache[e.memkey()] = e
		}
	}

	return es, nil
}

func (g *Goon) Run(q *datastore.Query) *Iterator {
	return &Iterator{
		g: g,
		i: q.Run(g.context),
	}
}

// Iterator is the result of running a query.
type Iterator struct {
	g *Goon
	i *datastore.Iterator
}

// Cursor returns a cursor for the iterator's current location.
func (t *Iterator) Cursor() (datastore.Cursor, error) {
	return t.i.Cursor()
}

// Next returns the entity of the next result. When there are no more results,
// datastore.Done is returned as the error. If dst is null (for a keys-only
// query), nil is returned as the entity.
//
// If the query is not keys only and dst is non-nil, it also loads the entity
// stored for that key into the struct pointer dst, with the same semantics
// and possible errors as for the Get function. This is returned as an Entity
// and cached in memory.
//
// If the query is keys only, dst must be passed as nil. Otherwise the cache
// will be populated with empty entities since there is no way to detect the
// case of a keys-only query.
//
// Refer to appengine/datastore.Iterator.Next:
// https://developers.google.com/appengine/docs/go/datastore/reference#Iterator.Next
func (t *Iterator) Next(dst interface{}) (*Entity, error) {
	k, err := t.i.Next(dst)
	if err != nil {
		return nil, err
	}

	if dst == nil {
		return nil, nil
	}

	e := NewEntity(k, dst)

	if !t.g.inTransaction {
		t.g.cache[e.memkey()] = e
	}

	return e, nil
}
<commit_msg>No longer polluting the cache with empty results from keys-only queries.<commit_after>/*
 * Copyright (c) 2012 Matt Jibson <matt.jibson@gmail.com>
 *
 * Permission to use, copy, modify, and distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */

package goon

import (
	"appengine/datastore"
	"reflect"
)

// Count returns the number of results for the query.
func (g *Goon) Count(q *datastore.Query) (int, error) {
	return q.Count(g.context)
}

// GetAll returns the entities in the given query. Returned entities (but not
// the query itself) are cached in local memory (but not memcache).
//
// Returns nil if q is a keys-only query or dst is not a slice.
//
// Refer to appengine/datastore.Query.Count:
// https://developers.google.com/appengine/docs/go/datastore/reference#Query.GetAll
//
// Example usage:
//   g := goon.NewGoon(r)
//   q := datastore.NewQuery("Group")
//   var gg []Group // (or []*Group)
//   es, err := g.GetAll(q, &gg)
func (g *Goon) GetAll(q *datastore.Query, dst interface{}) ([]*Entity, error) {
	keys, err := q.GetAll(g.context, dst)
	if err != nil {
		g.error(err)
		return nil, err
	}

	keysOnly := dst == nil

	var v reflect.Value
	var t reflect.Type
	if !keysOnly {
		v = reflect.Indirect(reflect.ValueOf(dst))
		t = v.Type()

		// try to detect a keys-only query
		if t.Kind() != reflect.Slice || v.Len() != len(keys) {
			keysOnly = true
		}
	}

	es := make([]*Entity, len(keys))

	for i, k := range keys {
		var e *Entity
		if keysOnly {
			e = NewEntity(k, nil)
		} else {
			e = NewEntity(k, v.Index(i).Interface())
		}
		es[i] = e

		// Do not pollute the cache with empty results from keys-only queries
		if !g.inTransaction && !keysOnly {
			g.cache[e.memkey()] = e
		}
	}

	return es, nil
}

func (g *Goon) Run(q *datastore.Query) *Iterator {
	return &Iterator{
		g: g,
		i: q.Run(g.context),
	}
}

// Iterator is the result of running a query.
type Iterator struct {
	g *Goon
	i *datastore.Iterator
}

// Cursor returns a cursor for the iterator's current location.
func (t *Iterator) Cursor() (datastore.Cursor, error) {
	return t.i.Cursor()
}

// Next returns the entity of the next result. When there are no more results,
// datastore.Done is returned as the error. If dst is null (for a keys-only
// query), nil is returned as the entity.
//
// If the query is not keys only and dst is non-nil, it also loads the entity
// stored for that key into the struct pointer dst, with the same semantics
// and possible errors as for the Get function. This is returned as an Entity
// and cached in memory.
//
// If the query is keys only, dst must be passed as nil. Otherwise the cache
// will be populated with empty entities since there is no way to detect the
// case of a keys-only query.
//
// Refer to appengine/datastore.Iterator.Next:
// https://developers.google.com/appengine/docs/go/datastore/reference#Iterator.Next
func (t *Iterator) Next(dst interface{}) (*Entity, error) {
	k, err := t.i.Next(dst)
	if err != nil {
		return nil, err
	}

	if dst == nil {
		return nil, nil
	}

	e := NewEntity(k, dst)

	if !t.g.inTransaction {
		t.g.cache[e.memkey()] = e
	}

	return e, nil
}
<|endoftext|>


<commit_before>package gosql

const (
	QueryKeyAnd string = "%and" // AND
	QueryKeyOr         = "%or"  // OR
)

func IsQueryKey1(str string) bool {
	switch str {
	case QueryKeyAnd, QueryKeyOr:
		return true
	}
	return false
}

const (
	QueryKeyEq      string = "%eq"   // 等于
	QueryKeyNe             = "%ne"   // 不等于
	QueryKeyLt             = "%lt"   // 小于
	QueryKeyLte            = "%lte"  // 小于等于
	QueryKeyGt             = "%gt"   // 大于
	QueryKeyGte            = "%gte"  // 大于等于
	QueryKeyLike           = "%like" // 模糊搜索
	QueryKeyIn             = "%in"   // TODO: 暂时不支持
	QueryKeyBetween        = "%bt"   // TODO: 暂时不支持
)

func IsQueryKey2(str string) bool {
	switch str {
	case QueryKeyEq, QueryKeyNe:
		return true
	case QueryKeyLt, QueryKeyLte:
		return true
	case QueryKeyGt, QueryKeyGte:
		return true
	case QueryKeyLike:
		return true
	}
	return false
}

type IQuery interface {
	IsAnonymous() bool
	QueryKey() string
	QueryValue() []IQuery
}

type QueryRoot struct {
	Value []IQuery
}

func (q *QueryRoot) IsAnonymous() bool {
	return false
}

func (q *QueryRoot) QueryKey() string {
	return QueryKeyAnd
}

func (q *QueryRoot) QueryValue() []IQuery {
	return q.Value
}

type QueryElem struct {
	anonymous bool
	Key       string
	Value     []IQuery
}

func (q *QueryElem) IsAnonymous() bool {
	return q.anonymous
}

func (q *QueryElem) QueryKey() string {
	return q.Key
}

func (q *QueryElem) QueryValue() []IQuery {
	return q.Value
}

type QueryValue struct {
	Key   string
	Field string
	Value interface{}
}

func (q *QueryValue) IsAnonymous() bool {
	return false
}

func (q *QueryValue) QueryKey() string {
	return q.Key
}

func (q *QueryValue) QueryValue() []IQuery {
	return nil
}
<commit_msg>feat: simplify query interface<commit_after>package gosql

const (
	QueryKeyAnd string = "%and" // AND
	QueryKeyOr         = "%or"  // OR
)

func IsQueryKey1(str string) bool {
	switch str {
	case QueryKeyAnd, QueryKeyOr:
		return true
	}
	return false
}

const (
	QueryKeyEq      string = "%eq"   // 等于
	QueryKeyNe             = "%ne"   // 不等于
	QueryKeyLt             = "%lt"   // 小于
	QueryKeyLte            = "%lte"  // 小于等于
	QueryKeyGt             = "%gt"   // 大于
	QueryKeyGte            = "%gte"  // 大于等于
	QueryKeyLike           = "%like" // 模糊搜索
	QueryKeyIn             = "%in"   // TODO: 暂时不支持
	QueryKeyBetween        = "%bt"   // TODO: 暂时不支持
)

func IsQueryKey2(str string) bool {
	switch str {
	case QueryKeyEq, QueryKeyNe:
		return true
	case QueryKeyLt, QueryKeyLte:
		return true
	case QueryKeyGt, QueryKeyGte:
		return true
	case QueryKeyLike:
		return true
	}
	return false
}

type IQuery interface {
	IsAnonymous() bool
}

type QueryRoot struct {
	Value []IQuery
}

func (q *QueryRoot) IsAnonymous() bool {
	return false
}

type QueryElem struct {
	anonymous bool
	Key       string
	Value     []IQuery
}

func (q *QueryElem) IsAnonymous() bool {
	return q.anonymous
}

type QueryValue struct {
	Key   string
	Field string
	Value interface{}
}

func (q *QueryValue) IsAnonymous() bool {
	return false
}
<|endoftext|>


<commit_before>package job

import (
	"context"
	"github.com/pkg/errors"
	"github.com/problame/go-streamrpc"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/zrepl/zrepl/config"
	"github.com/zrepl/zrepl/daemon/job/reset"
	"github.com/zrepl/zrepl/daemon/job/wakeup"
	"github.com/zrepl/zrepl/daemon/transport/connecter"
	"github.com/zrepl/zrepl/daemon/filters"
	"github.com/zrepl/zrepl/daemon/pruner"
	"github.com/zrepl/zrepl/endpoint"
	"github.com/zrepl/zrepl/replication"
	"github.com/zrepl/zrepl/zfs"
	"sync"
	"github.com/zrepl/zrepl/daemon/logging"
	"github.com/zrepl/zrepl/daemon/snapper"
	"time"
)

type ActiveSide struct {
	mode          activeMode
	name          string
	clientFactory *connecter.ClientFactory

	prunerFactory *pruner.PrunerFactory

	promRepStateSecs *prometheus.HistogramVec // labels: state
	promPruneSecs *prometheus.HistogramVec // labels: prune_side
	promBytesReplicated *prometheus.CounterVec // labels: filesystem

	tasksMtx sync.Mutex
	tasks    activeSideTasks
}

type activeSideTasks struct {
	replication *replication.Replication
	prunerSender, prunerReceiver *pruner.Pruner
}

func (a *ActiveSide) updateTasks(u func(*activeSideTasks)) activeSideTasks {
	a.tasksMtx.Lock()
	defer a.tasksMtx.Unlock()
	var copy activeSideTasks
	copy = a.tasks
	if u == nil {
		return copy
	}
	u(&copy)
	a.tasks = copy
	return copy
}

type activeMode interface {
	SenderReceiver(client *streamrpc.Client) (replication.Sender, replication.Receiver, error)
	Type() Type
	RunPeriodic(ctx context.Context, wakeUpCommon chan<- struct{})
}

type modePush struct {
	fsfilter         endpoint.FSFilter
	snapper *snapper.PeriodicOrManual
}

func (m *modePush) SenderReceiver(client *streamrpc.Client) (replication.Sender, replication.Receiver, error) {
	sender := endpoint.NewSender(m.fsfilter)
	receiver := endpoint.NewRemote(client)
	return sender, receiver, nil
}

func (m *modePush) Type() Type { return TypePush }

func (m *modePush) RunPeriodic(ctx context.Context, wakeUpCommon chan <- struct{}) {
	m.snapper.Run(ctx, wakeUpCommon)
}

func modePushFromConfig(g *config.Global, in *config.PushJob) (*modePush, error) {
	m := &modePush{}
	fsf, err := filters.DatasetMapFilterFromConfig(in.Filesystems)
	if err != nil {
		return nil, errors.Wrap(err, "cannnot build filesystem filter")
	}
	m.fsfilter = fsf

	if m.snapper, err = snapper.FromConfig(g, fsf, in.Snapshotting); err != nil {
		return nil, errors.Wrap(err, "cannot build snapper")
	}

	return m, nil
}

type modePull struct {
	rootFS   *zfs.DatasetPath
	interval time.Duration
}

func (m *modePull) SenderReceiver(client *streamrpc.Client) (replication.Sender, replication.Receiver, error) {
	sender := endpoint.NewRemote(client)
	receiver, err := endpoint.NewReceiver(m.rootFS)
	return sender, receiver, err
}

func (*modePull) Type() Type { return TypePull }

func (m *modePull) RunPeriodic(ctx context.Context, wakeUpCommon chan<- struct{}) {
	t := time.NewTicker(m.interval)
	defer t.Stop()
	for {
		select {
		case <-t.C:
			select {
			case wakeUpCommon <- struct{}{}:
			default:
				GetLogger(ctx).
					WithField("pull_interval", m.interval).
					Warn("pull job took longer than pull interval")
				wakeUpCommon <- struct{}{} // block anyways, to queue up the wakeup
			}
		case <-ctx.Done():
			return
		}
	}
}

func modePullFromConfig(g *config.Global, in *config.PullJob) (m *modePull, err error) {
	m = &modePull{}
	if in.Interval <= 0 {
		return nil, errors.New("interval must be positive")
	}
	m.interval = in.Interval

	m.rootFS, err = zfs.NewDatasetPath(in.RootFS)
	if err != nil {
		return nil, errors.New("RootFS is not a valid zfs filesystem path")
	}
	if m.rootFS.Length() <= 0 {
		return nil, errors.New("RootFS must not be empty") // duplicates error check of receiver
	}

	return m, nil
}

func activeSide(g *config.Global, in *config.ActiveJob, mode activeMode) (j *ActiveSide, err error) {

	j = &ActiveSide{mode: mode}
	j.name = in.Name
	j.promRepStateSecs = prometheus.NewHistogramVec(prometheus.HistogramOpts{
		Namespace:   "zrepl",
		Subsystem:   "replication",
		Name:        "state_time",
		Help:        "seconds spent during replication",
		ConstLabels: prometheus.Labels{"zrepl_job":j.name},
	}, []string{"state"})
	j.promBytesReplicated = prometheus.NewCounterVec(prometheus.CounterOpts{
		Namespace:   "zrepl",
		Subsystem:   "replication",
		Name:        "bytes_replicated",
		Help:        "number of bytes replicated from sender to receiver per filesystem",
		ConstLabels: prometheus.Labels{"zrepl_job":j.name},
	}, []string{"filesystem"})

	j.clientFactory, err = connecter.FromConfig(g, in.Connect)
	if err != nil {
		return nil, errors.Wrap(err, "cannot build client")
	}

	j.promPruneSecs = prometheus.NewHistogramVec(prometheus.HistogramOpts{
		Namespace:   "zrepl",
		Subsystem:   "pruning",
		Name:        "time",
		Help:        "seconds spent in pruner",
		ConstLabels: prometheus.Labels{"zrepl_job":j.name},
	}, []string{"prune_side"})
	j.prunerFactory, err = pruner.NewPrunerFactory(in.Pruning, j.promPruneSecs)
	if err != nil {
		return nil, err
	}

	return j, nil
}

func (j *ActiveSide) RegisterMetrics(registerer prometheus.Registerer) {
	registerer.MustRegister(j.promRepStateSecs)
	registerer.MustRegister(j.promPruneSecs)
	registerer.MustRegister(j.promBytesReplicated)
}

func (j *ActiveSide) Name() string { return j.name }

type ActiveSideStatus struct {
	Replication *replication.Report
	PruningSender, PruningReceiver *pruner.Report
}

func (j *ActiveSide) Status() *Status {
	tasks := j.updateTasks(nil)

	s := &ActiveSideStatus{}
	t := j.mode.Type()
	if tasks.replication != nil {
		s.Replication = tasks.replication.Report()
	}
	if tasks.prunerSender != nil {
		s.PruningSender = tasks.prunerSender.Report()
	}
	if tasks.prunerReceiver != nil {
		s.PruningReceiver = tasks.prunerReceiver.Report()
	}
	return &Status{Type: t, JobSpecific: s}
}

func (j *ActiveSide) Run(ctx context.Context) {
	log := GetLogger(ctx)
	ctx = logging.WithSubsystemLoggers(ctx, log)

	defer log.Info("job exiting")

	periodicDone := make(chan struct{})
	ctx, cancel := context.WithCancel(ctx)
	defer cancel()
	go j.mode.RunPeriodic(ctx, periodicDone)

	invocationCount := 0
outer:
	for {
		log.Info("wait for wakeups")
		select {
		case <-ctx.Done():
			log.WithError(ctx.Err()).Info("context")
			break outer

		case <-wakeup.Wait(ctx):
		case <-periodicDone:
		}
		invocationCount++
		invLog := log.WithField("invocation", invocationCount)
		j.do(WithLogger(ctx, invLog))
	}
}

func (j *ActiveSide) do(ctx context.Context) {

	log := GetLogger(ctx)
	ctx = logging.WithSubsystemLoggers(ctx, log)

	// allow cancellation of an invocation (this function)
	ctx, cancelThisRun := context.WithCancel(ctx)
	defer cancelThisRun()
	runDone := make(chan struct{})
	defer close(runDone)
	go func() {
		select {
		case <-runDone:
		case <-reset.Wait(ctx):
			log.Info("reset received, cancelling current invocation")
			cancelThisRun()
		case <-ctx.Done():
		}
	}()

	client, err := j.clientFactory.NewClient()
	if err != nil {
		log.WithError(err).Error("factory cannot instantiate streamrpc client")
	}
	defer client.Close(ctx)

	sender, receiver, err := j.mode.SenderReceiver(client)

	tasks := j.updateTasks(func(tasks *activeSideTasks) {
		// reset it
		*tasks = activeSideTasks{}
		tasks.replication = replication.NewReplication(j.promRepStateSecs, j.promBytesReplicated)
	})

	log.Info("start replication")
	tasks.replication.Drive(ctx, sender, receiver)

	tasks = j.updateTasks(func(tasks *activeSideTasks) {
		tasks.prunerSender = j.prunerFactory.BuildSenderPruner(ctx, sender, sender)
		tasks.prunerReceiver = j.prunerFactory.BuildReceiverPruner(ctx, receiver, sender)
	})
	log.Info("start pruning sender")
	tasks.prunerSender.Prune()
	log.Info("start pruning receiver")
	tasks.prunerReceiver.Prune()
}
<commit_msg>daemon/job/active: push mode: awful hack for handling of concurrent snapshots + stale remote operation<commit_after>package job

import (
	"context"
	"github.com/pkg/errors"
	"github.com/problame/go-streamrpc"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/zrepl/zrepl/config"
	"github.com/zrepl/zrepl/daemon/job/reset"
	"github.com/zrepl/zrepl/daemon/job/wakeup"
	"github.com/zrepl/zrepl/daemon/transport/connecter"
	"github.com/zrepl/zrepl/daemon/filters"
	"github.com/zrepl/zrepl/daemon/pruner"
	"github.com/zrepl/zrepl/endpoint"
	"github.com/zrepl/zrepl/replication"
	"github.com/zrepl/zrepl/zfs"
	"sync"
	"github.com/zrepl/zrepl/daemon/logging"
	"github.com/zrepl/zrepl/daemon/snapper"
	"time"
)

type ActiveSide struct {
	mode          activeMode
	name          string
	clientFactory *connecter.ClientFactory

	prunerFactory *pruner.PrunerFactory

	promRepStateSecs *prometheus.HistogramVec // labels: state
	promPruneSecs *prometheus.HistogramVec // labels: prune_side
	promBytesReplicated *prometheus.CounterVec // labels: filesystem

	tasksMtx sync.Mutex
	tasks    activeSideTasks
}

type activeSideTasks struct {
	replication *replication.Replication
	prunerSender, prunerReceiver *pruner.Pruner
}

func (a *ActiveSide) updateTasks(u func(*activeSideTasks)) activeSideTasks {
	a.tasksMtx.Lock()
	defer a.tasksMtx.Unlock()
	var copy activeSideTasks
	copy = a.tasks
	if u == nil {
		return copy
	}
	u(&copy)
	a.tasks = copy
	return copy
}

type activeMode interface {
	SenderReceiver(client *streamrpc.Client) (replication.Sender, replication.Receiver, error)
	Type() Type
	RunPeriodic(ctx context.Context, wakeUpCommon chan<- struct{})
}

type modePush struct {
	fsfilter         endpoint.FSFilter
	snapper *snapper.PeriodicOrManual
}

func (m *modePush) SenderReceiver(client *streamrpc.Client) (replication.Sender, replication.Receiver, error) {
	sender := endpoint.NewSender(m.fsfilter)
	receiver := endpoint.NewRemote(client)
	return sender, receiver, nil
}

func (m *modePush) Type() Type { return TypePush }

func (m *modePush) RunPeriodic(ctx context.Context, wakeUpCommon chan <- struct{}) {
	m.snapper.Run(ctx, wakeUpCommon)
}

func modePushFromConfig(g *config.Global, in *config.PushJob) (*modePush, error) {
	m := &modePush{}
	fsf, err := filters.DatasetMapFilterFromConfig(in.Filesystems)
	if err != nil {
		return nil, errors.Wrap(err, "cannnot build filesystem filter")
	}
	m.fsfilter = fsf

	if m.snapper, err = snapper.FromConfig(g, fsf, in.Snapshotting); err != nil {
		return nil, errors.Wrap(err, "cannot build snapper")
	}

	return m, nil
}

type modePull struct {
	rootFS   *zfs.DatasetPath
	interval time.Duration
}

func (m *modePull) SenderReceiver(client *streamrpc.Client) (replication.Sender, replication.Receiver, error) {
	sender := endpoint.NewRemote(client)
	receiver, err := endpoint.NewReceiver(m.rootFS)
	return sender, receiver, err
}

func (*modePull) Type() Type { return TypePull }

func (m *modePull) RunPeriodic(ctx context.Context, wakeUpCommon chan<- struct{}) {
	t := time.NewTicker(m.interval)
	defer t.Stop()
	for {
		select {
		case <-t.C:
			select {
			case wakeUpCommon <- struct{}{}:
			default:
				GetLogger(ctx).
					WithField("pull_interval", m.interval).
					Warn("pull job took longer than pull interval")
				wakeUpCommon <- struct{}{} // block anyways, to queue up the wakeup
			}
		case <-ctx.Done():
			return
		}
	}
}

func modePullFromConfig(g *config.Global, in *config.PullJob) (m *modePull, err error) {
	m = &modePull{}
	if in.Interval <= 0 {
		return nil, errors.New("interval must be positive")
	}
	m.interval = in.Interval

	m.rootFS, err = zfs.NewDatasetPath(in.RootFS)
	if err != nil {
		return nil, errors.New("RootFS is not a valid zfs filesystem path")
	}
	if m.rootFS.Length() <= 0 {
		return nil, errors.New("RootFS must not be empty") // duplicates error check of receiver
	}

	return m, nil
}

func activeSide(g *config.Global, in *config.ActiveJob, mode activeMode) (j *ActiveSide, err error) {

	j = &ActiveSide{mode: mode}
	j.name = in.Name
	j.promRepStateSecs = prometheus.NewHistogramVec(prometheus.HistogramOpts{
		Namespace:   "zrepl",
		Subsystem:   "replication",
		Name:        "state_time",
		Help:        "seconds spent during replication",
		ConstLabels: prometheus.Labels{"zrepl_job":j.name},
	}, []string{"state"})
	j.promBytesReplicated = prometheus.NewCounterVec(prometheus.CounterOpts{
		Namespace:   "zrepl",
		Subsystem:   "replication",
		Name:        "bytes_replicated",
		Help:        "number of bytes replicated from sender to receiver per filesystem",
		ConstLabels: prometheus.Labels{"zrepl_job":j.name},
	}, []string{"filesystem"})

	j.clientFactory, err = connecter.FromConfig(g, in.Connect)
	if err != nil {
		return nil, errors.Wrap(err, "cannot build client")
	}

	j.promPruneSecs = prometheus.NewHistogramVec(prometheus.HistogramOpts{
		Namespace:   "zrepl",
		Subsystem:   "pruning",
		Name:        "time",
		Help:        "seconds spent in pruner",
		ConstLabels: prometheus.Labels{"zrepl_job":j.name},
	}, []string{"prune_side"})
	j.prunerFactory, err = pruner.NewPrunerFactory(in.Pruning, j.promPruneSecs)
	if err != nil {
		return nil, err
	}

	return j, nil
}

func (j *ActiveSide) RegisterMetrics(registerer prometheus.Registerer) {
	registerer.MustRegister(j.promRepStateSecs)
	registerer.MustRegister(j.promPruneSecs)
	registerer.MustRegister(j.promBytesReplicated)
}

func (j *ActiveSide) Name() string { return j.name }

type ActiveSideStatus struct {
	Replication *replication.Report
	PruningSender, PruningReceiver *pruner.Report
}

func (j *ActiveSide) Status() *Status {
	tasks := j.updateTasks(nil)

	s := &ActiveSideStatus{}
	t := j.mode.Type()
	if tasks.replication != nil {
		s.Replication = tasks.replication.Report()
	}
	if tasks.prunerSender != nil {
		s.PruningSender = tasks.prunerSender.Report()
	}
	if tasks.prunerReceiver != nil {
		s.PruningReceiver = tasks.prunerReceiver.Report()
	}
	return &Status{Type: t, JobSpecific: s}
}

func (j *ActiveSide) Run(ctx context.Context) {
	log := GetLogger(ctx)
	ctx = logging.WithSubsystemLoggers(ctx, log)

	defer log.Info("job exiting")

	periodicDone := make(chan struct{})
	ctx, cancel := context.WithCancel(ctx)
	defer cancel()
	go j.mode.RunPeriodic(ctx, periodicDone)

	invocationCount := 0
outer:
	for {
		log.Info("wait for wakeups")
		select {
		case <-ctx.Done():
			log.WithError(ctx.Err()).Info("context")
			break outer

		case <-wakeup.Wait(ctx):
		case <-periodicDone:
		}
		invocationCount++
		invLog := log.WithField("invocation", invocationCount)
		j.do(WithLogger(ctx, invLog), periodicDone)
	}
}

func (j *ActiveSide) do(ctx context.Context, periodicWakeup <-chan struct{}) {

	log := GetLogger(ctx)
	ctx = logging.WithSubsystemLoggers(ctx, log)

	// allow cancellation of an invocation (this function)
	ctx, cancelThisRun := context.WithCancel(ctx)
	defer cancelThisRun()
	runDone := make(chan struct{})
	defer close(runDone)
	go func() {
		select {
		case <-runDone:
		case <-reset.Wait(ctx):
			log.Info("reset received, cancelling current invocation")
			cancelThisRun()
		case <-ctx.Done():
		}
	}()

	client, err := j.clientFactory.NewClient()
	if err != nil {
		log.WithError(err).Error("factory cannot instantiate streamrpc client")
	}
	defer client.Close(ctx)

	sender, receiver, err := j.mode.SenderReceiver(client)

	tasks := j.updateTasks(func(tasks *activeSideTasks) {
		// reset it
		*tasks = activeSideTasks{}
		tasks.replication = replication.NewReplication(j.promRepStateSecs, j.promBytesReplicated)
	})

	log.Info("start replication")
	replicationDone := make(chan struct{})
	replicationCtx, replicationCancel := context.WithCancel(ctx)
	defer replicationCancel()
	go func() {
		tasks.replication.Drive(replicationCtx, sender, receiver)
		close(replicationDone)
	}()
	outer:
	for {
		select {
		case <-replicationDone:
			// fine!
			break outer
		case <-periodicWakeup:
			// Replication took longer than the periodic interval.
			//
			// For pull jobs, this isn't so bad because nothing changes on the active side
			// if replication doesn't go forward.
			//
			// For push jobs, this means snapshots were taken.
			// We need to invoke the pruner now, because otherwise an infinitely stuck replication
			// will cause this side to fill up with snapshots.
			//
			// However, there are cases where replication progresses and just takes longer,
			// and we don't want these situations be interrupted by a prune, which will require
			// re-planning and starting over (think of initial replication as an example).
			//
			// Therefore, we prohibit pruning of snapshots that are part of the current replication plan.
			// If there is no such plan, we kill the replication.

			if j.mode.Type() == TypePush {

				rep := tasks.replication.Report()
				state, err := replication.StateString(rep.Status)
				if err != nil {
					panic(err)
				}

				switch state {
				case replication.Planning:
					fallthrough
				case replication.PlanningError:
					fallthrough
				case replication.WorkingWait:
					log.WithField("repl_state", state.String()).
						Info("cancelling replication after new snapshots invalidated its current state")
					replicationCancel()
					log.Info("waiting for replication to stop")
					<-replicationDone // no need to wait for ctx.Done, replication is already bound to global cancel
					break outer
				default:
					log.WithField("repl_state", state.String()).
						Warn("new snapshots while replication is running and making progress")
				}

			}

		}
	}

	var pruningWg sync.WaitGroup
	log.Info("start pruning sender")
	pruningWg.Add(1)
	go func() {
		defer pruningWg.Done()
		tasks := j.updateTasks(func(tasks *activeSideTasks) {
			tasks.prunerSender = j.prunerFactory.BuildSenderPruner(ctx, sender, sender)
		})
		tasks.prunerSender.Prune()
		// FIXME no need to do the cancellation dance with sender, we know it's local for push
		// FIXME and we don't worry about pull ATM
	}()
	log.Info("start pruning receiver")
	pruningWg.Add(1)
	go func() {
		defer pruningWg.Done()

		receiverPrunerCtx, receiverPrunerCancel := context.WithCancel(ctx)
		defer receiverPrunerCancel()
		tasks := j.updateTasks(func(tasks *activeSideTasks) {
			tasks.prunerReceiver = j.prunerFactory.BuildReceiverPruner(receiverPrunerCtx, receiver, sender)
		})
		receiverPrunerDone := make(chan struct{})
		go func() {
			defer close(receiverPrunerDone)
			tasks.prunerReceiver.Prune()
		}()

		outer:
		for {
			select {
			case <-receiverPrunerDone:
				// fine!
				break outer
			case <-periodicWakeup:
				// see comments for similar apporach with replication above
				if j.mode.Type() == TypePush {
					rep := tasks.prunerReceiver.Report()
					state, err := pruner.StateString(rep.State)
					if err != nil {
						panic(err)
					}
					switch state {
					case pruner.PlanWait:
						fallthrough
					case pruner.ExecWait:
						log.WithField("pruner_state", state.String()).
							Info("cancelling failing prune on receiver because new snapshots were taken on sender")
						receiverPrunerCancel()
						log.Info("waiting for receiver pruner to stop")
						<-receiverPrunerDone
						break outer
					default:
						log.WithField("pruner_state", state.String()).
							Warn("new snapshots while prune on receiver is still running")
					}
				}
			}
		}

	}()

	pruningWg.Wait() // if pruners handle ctx cancellation correctly, we don't need to wait for it here

}
<|endoftext|>