Initial QSfera import
This commit is contained in:
+44
@@ -0,0 +1,44 @@
|
||||
// Copyright 2018 The OPA Authors. All rights reserved.
|
||||
// Use of this source code is governed by an Apache2
|
||||
// license that can be found in the LICENSE file.
|
||||
|
||||
package util
|
||||
|
||||
import (
|
||||
"math/rand"
|
||||
"time"
|
||||
)
|
||||
|
||||
// DefaultBackoff returns a delay with an exponential backoff based on the
|
||||
// number of retries.
|
||||
func DefaultBackoff(base, maxNS float64, retries int) time.Duration {
|
||||
return Backoff(base, maxNS, .2, 1.6, retries)
|
||||
}
|
||||
|
||||
// Backoff returns a delay with an exponential backoff based on the number of
|
||||
// retries. Same algorithm used in gRPC.
|
||||
// Note that if maxNS is smaller than base, the backoff will still be capped at
|
||||
// maxNS.
|
||||
func Backoff(base, maxNS, jitter, factor float64, retries int) time.Duration {
|
||||
if retries == 0 {
|
||||
return 0
|
||||
}
|
||||
|
||||
backoff, maxNS := base, maxNS
|
||||
for backoff < maxNS && retries > 0 {
|
||||
backoff *= factor
|
||||
retries--
|
||||
}
|
||||
if backoff > maxNS {
|
||||
backoff = maxNS
|
||||
}
|
||||
|
||||
// Randomize backoff delays so that if a cluster of requests start at
|
||||
// the same time, they won't operate in lockstep.
|
||||
backoff *= 1 + jitter*(rand.Float64()*2-1)
|
||||
if backoff < 0 {
|
||||
return 0
|
||||
}
|
||||
|
||||
return time.Duration(backoff)
|
||||
}
|
||||
+32
@@ -0,0 +1,32 @@
|
||||
package util
|
||||
|
||||
import (
|
||||
"github.com/open-policy-agent/opa/v1/metrics"
|
||||
)
|
||||
|
||||
// This prevents getting blocked forever writing to a full buffer, in case another routine fills the last space.
|
||||
// Retrying maxEventRetry times to drop the oldest event. Dropping the incoming event if there still isn't room.
|
||||
const maxEventRetry = 1000
|
||||
|
||||
// PushFIFO pushes data into a buffered channel without blocking when full, making room by dropping the oldest data.
|
||||
// An optional metric can be recorded when data is dropped.
|
||||
func PushFIFO[T any](buffer chan T, data T, metrics metrics.Metrics, metricName string) {
|
||||
|
||||
for range maxEventRetry {
|
||||
// non-blocking send to the buffer, to prevent blocking if buffer is full so room can be made.
|
||||
select {
|
||||
case buffer <- data:
|
||||
return
|
||||
default:
|
||||
}
|
||||
|
||||
// non-blocking drop from the buffer to make room for incoming event
|
||||
select {
|
||||
case <-buffer:
|
||||
if metrics != nil && metricName != "" {
|
||||
metrics.Counter(metricName).Incr()
|
||||
}
|
||||
default:
|
||||
}
|
||||
}
|
||||
}
|
||||
+22
@@ -0,0 +1,22 @@
|
||||
// Copyright 2018 The OPA Authors. All rights reserved.
|
||||
// Use of this source code is governed by an Apache2
|
||||
// license that can be found in the LICENSE file.
|
||||
|
||||
package util
|
||||
|
||||
import (
|
||||
"io"
|
||||
"net/http"
|
||||
)
|
||||
|
||||
// Close reads the remaining bytes from the response and then closes it to
|
||||
// ensure that the connection is freed. If the body is not read and closed, a
|
||||
// leak can occur.
|
||||
func Close(resp *http.Response) {
|
||||
if resp != nil && resp.Body != nil {
|
||||
if _, err := io.Copy(io.Discard, resp.Body); err != nil {
|
||||
return
|
||||
}
|
||||
resp.Body.Close()
|
||||
}
|
||||
}
|
||||
+169
@@ -0,0 +1,169 @@
|
||||
// Copyright 2016 The OPA Authors. All rights reserved.
|
||||
// Use of this source code is governed by an Apache2
|
||||
// license that can be found in the LICENSE file.
|
||||
|
||||
package util
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"math/big"
|
||||
)
|
||||
|
||||
// Compare returns 0 if a equals b, -1 if a is less than b, and 1 if b is than a.
|
||||
//
|
||||
// For comparison between values of different types, the following ordering is used:
|
||||
// nil < bool < int, float64 < string < []any < map[string]any. Slices and maps
|
||||
// are compared recursively. If one slice or map is a subset of the other slice or map
|
||||
// it is considered "less than". Nil is always equal to nil.
|
||||
func Compare(a, b any) int {
|
||||
aSortOrder := sortOrder(a)
|
||||
bSortOrder := sortOrder(b)
|
||||
if aSortOrder < bSortOrder {
|
||||
return -1
|
||||
} else if bSortOrder < aSortOrder {
|
||||
return 1
|
||||
}
|
||||
switch a := a.(type) {
|
||||
case nil:
|
||||
return 0
|
||||
case bool:
|
||||
switch b := b.(type) {
|
||||
case bool:
|
||||
if a == b {
|
||||
return 0
|
||||
}
|
||||
if !a {
|
||||
return -1
|
||||
}
|
||||
return 1
|
||||
}
|
||||
case json.Number:
|
||||
switch b := b.(type) {
|
||||
case json.Number:
|
||||
return compareJSONNumber(a, b)
|
||||
}
|
||||
case int:
|
||||
switch b := b.(type) {
|
||||
case int:
|
||||
if a == b {
|
||||
return 0
|
||||
} else if a < b {
|
||||
return -1
|
||||
}
|
||||
return 1
|
||||
}
|
||||
case float64:
|
||||
switch b := b.(type) {
|
||||
case float64:
|
||||
if a == b {
|
||||
return 0
|
||||
} else if a < b {
|
||||
return -1
|
||||
}
|
||||
return 1
|
||||
}
|
||||
case string:
|
||||
switch b := b.(type) {
|
||||
case string:
|
||||
if a == b {
|
||||
return 0
|
||||
} else if a < b {
|
||||
return -1
|
||||
}
|
||||
return 1
|
||||
}
|
||||
case []any:
|
||||
switch b := b.(type) {
|
||||
case []any:
|
||||
bLen := len(b)
|
||||
aLen := len(a)
|
||||
minLen := min(bLen, aLen)
|
||||
for i := range minLen {
|
||||
cmp := Compare(a[i], b[i])
|
||||
if cmp != 0 {
|
||||
return cmp
|
||||
}
|
||||
}
|
||||
if aLen == bLen {
|
||||
return 0
|
||||
} else if aLen < bLen {
|
||||
return -1
|
||||
}
|
||||
return 1
|
||||
}
|
||||
case map[string]any:
|
||||
switch b := b.(type) {
|
||||
case map[string]any:
|
||||
aKeys := KeysSorted(a)
|
||||
bKeys := KeysSorted(b)
|
||||
aLen := len(aKeys)
|
||||
bLen := len(bKeys)
|
||||
minLen := min(bLen, aLen)
|
||||
for i := range minLen {
|
||||
if aKeys[i] < bKeys[i] {
|
||||
return -1
|
||||
} else if bKeys[i] < aKeys[i] {
|
||||
return 1
|
||||
}
|
||||
aVal := a[aKeys[i]]
|
||||
bVal := b[bKeys[i]]
|
||||
cmp := Compare(aVal, bVal)
|
||||
if cmp != 0 {
|
||||
return cmp
|
||||
}
|
||||
}
|
||||
if aLen == bLen {
|
||||
return 0
|
||||
} else if aLen < bLen {
|
||||
return -1
|
||||
}
|
||||
return 1
|
||||
}
|
||||
}
|
||||
|
||||
panic(fmt.Sprintf("illegal arguments of type %T and type %T", a, b))
|
||||
}
|
||||
|
||||
const (
|
||||
nilSort = iota
|
||||
boolSort = iota
|
||||
numberSort = iota
|
||||
stringSort = iota
|
||||
arraySort = iota
|
||||
objectSort = iota
|
||||
)
|
||||
|
||||
func compareJSONNumber(a, b json.Number) int {
|
||||
bigA, ok := new(big.Float).SetString(string(a))
|
||||
if !ok {
|
||||
panic("illegal value")
|
||||
}
|
||||
bigB, ok := new(big.Float).SetString(string(b))
|
||||
if !ok {
|
||||
panic("illegal value")
|
||||
}
|
||||
return bigA.Cmp(bigB)
|
||||
}
|
||||
|
||||
func sortOrder(v any) int {
|
||||
switch v.(type) {
|
||||
case nil:
|
||||
return nilSort
|
||||
case bool:
|
||||
return boolSort
|
||||
case json.Number:
|
||||
return numberSort
|
||||
case int:
|
||||
return numberSort
|
||||
case float64:
|
||||
return numberSort
|
||||
case string:
|
||||
return stringSort
|
||||
case []any:
|
||||
return arraySort
|
||||
case map[string]any:
|
||||
return objectSort
|
||||
}
|
||||
panic(fmt.Sprintf("illegal argument of type %T", v))
|
||||
}
|
||||
+31
@@ -0,0 +1,31 @@
|
||||
package decoding
|
||||
|
||||
import "context"
|
||||
|
||||
type requestContextKey string
|
||||
|
||||
// Note(philipc): We can add functions later to add the max request body length
|
||||
// to contexts, if we ever need to.
|
||||
const (
|
||||
reqCtxKeyMaxLen = requestContextKey("server-decoding-plugin-context-max-length")
|
||||
reqCtxKeyGzipMaxLen = requestContextKey("server-decoding-plugin-context-gzip-max-length")
|
||||
)
|
||||
|
||||
func AddServerDecodingMaxLen(ctx context.Context, maxLen int64) context.Context {
|
||||
return context.WithValue(ctx, reqCtxKeyMaxLen, maxLen)
|
||||
}
|
||||
|
||||
func AddServerDecodingGzipMaxLen(ctx context.Context, maxLen int64) context.Context {
|
||||
return context.WithValue(ctx, reqCtxKeyGzipMaxLen, maxLen)
|
||||
}
|
||||
|
||||
// Used for enforcing max body content limits when dealing with chunked requests.
|
||||
func GetServerDecodingMaxLen(ctx context.Context) (int64, bool) {
|
||||
maxLength, ok := ctx.Value(reqCtxKeyMaxLen).(int64)
|
||||
return maxLength, ok
|
||||
}
|
||||
|
||||
func GetServerDecodingGzipMaxLen(ctx context.Context) (int64, bool) {
|
||||
gzipMaxLength, ok := ctx.Value(reqCtxKeyGzipMaxLen).(int64)
|
||||
return gzipMaxLength, ok
|
||||
}
|
||||
+6
@@ -0,0 +1,6 @@
|
||||
// Copyright 2016 The OPA Authors. All rights reserved.
|
||||
// Use of this source code is governed by an Apache2
|
||||
// license that can be found in the LICENSE file.
|
||||
|
||||
// Package util provides generic utilities used throughout the policy engine.
|
||||
package util
|
||||
+59
@@ -0,0 +1,59 @@
|
||||
// Copyright 2017 The OPA Authors. All rights reserved.
|
||||
// Use of this source code is governed by an Apache2
|
||||
// license that can be found in the LICENSE file.
|
||||
|
||||
package util
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// EnumFlag implements the pflag.Value interface to provide enumerated command
|
||||
// line parameter values.
|
||||
type EnumFlag struct {
|
||||
defaultValue string
|
||||
vs []string
|
||||
i int
|
||||
}
|
||||
|
||||
// NewEnumFlag returns a new EnumFlag that has a defaultValue and vs enumerated
|
||||
// values.
|
||||
func NewEnumFlag(defaultValue string, vs []string) *EnumFlag {
|
||||
f := &EnumFlag{
|
||||
i: -1,
|
||||
vs: vs,
|
||||
defaultValue: defaultValue,
|
||||
}
|
||||
return f
|
||||
}
|
||||
|
||||
// Type returns the valid enumeration values.
|
||||
func (f *EnumFlag) Type() string {
|
||||
return "{" + strings.Join(f.vs, ",") + "}"
|
||||
}
|
||||
|
||||
// String returns the EnumValue's value as string.
|
||||
func (f *EnumFlag) String() string {
|
||||
if f.i == -1 {
|
||||
return f.defaultValue
|
||||
}
|
||||
return f.vs[f.i]
|
||||
}
|
||||
|
||||
// IsSet will return true if the EnumFlag has been set.
|
||||
func (f *EnumFlag) IsSet() bool {
|
||||
return f.i != -1
|
||||
}
|
||||
|
||||
// Set sets the enum value. If s is not a valid enum value, an error is
|
||||
// returned.
|
||||
func (f *EnumFlag) Set(s string) error {
|
||||
for i := range f.vs {
|
||||
if f.vs[i] == s {
|
||||
f.i = i
|
||||
return nil
|
||||
}
|
||||
}
|
||||
return fmt.Errorf("must be one of %v", f.Type())
|
||||
}
|
||||
+87
@@ -0,0 +1,87 @@
|
||||
// Copyright 2016 The OPA Authors. All rights reserved.
|
||||
// Use of this source code is governed by an Apache2
|
||||
// license that can be found in the LICENSE file.
|
||||
|
||||
package util
|
||||
|
||||
// Traversal defines a basic interface to perform traversals.
|
||||
type Traversal interface {
|
||||
|
||||
// Edges should return the neighbours of node "u".
|
||||
Edges(u T) []T
|
||||
|
||||
// Visited should return true if node "u" has already been visited in this
|
||||
// traversal. If the same traversal is used multiple times, the state that
|
||||
// tracks visited nodes should be reset.
|
||||
Visited(u T) bool
|
||||
}
|
||||
|
||||
// Equals should return true if node "u" equals node "v".
|
||||
type Equals func(u T, v T) bool
|
||||
|
||||
// Iter should return true to indicate stop.
|
||||
type Iter func(u T) bool
|
||||
|
||||
// DFS performs a depth first traversal calling f for each node starting from u.
|
||||
// If f returns true, traversal stops and DFS returns true.
|
||||
func DFS(t Traversal, f Iter, u T) bool {
|
||||
lifo := NewLIFO(u)
|
||||
for lifo.Size() > 0 {
|
||||
next, _ := lifo.Pop()
|
||||
if t.Visited(next) {
|
||||
continue
|
||||
}
|
||||
if f(next) {
|
||||
return true
|
||||
}
|
||||
for _, v := range t.Edges(next) {
|
||||
lifo.Push(v)
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// BFS performs a breadth first traversal calling f for each node starting from
|
||||
// u. If f returns true, traversal stops and BFS returns true.
|
||||
func BFS(t Traversal, f Iter, u T) bool {
|
||||
fifo := NewFIFO(u)
|
||||
for fifo.Size() > 0 {
|
||||
next, _ := fifo.Pop()
|
||||
if t.Visited(next) {
|
||||
continue
|
||||
}
|
||||
if f(next) {
|
||||
return true
|
||||
}
|
||||
for _, v := range t.Edges(next) {
|
||||
fifo.Push(v)
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// DFSPath returns a path from node a to node z found by performing
|
||||
// a depth first traversal. If no path is found, an empty slice is returned.
|
||||
func DFSPath(t Traversal, eq Equals, a, z T) []T {
|
||||
p := dfsRecursive(t, eq, a, z, []T{})
|
||||
for i := len(p)/2 - 1; i >= 0; i-- {
|
||||
o := len(p) - i - 1
|
||||
p[i], p[o] = p[o], p[i]
|
||||
}
|
||||
return p
|
||||
}
|
||||
|
||||
func dfsRecursive(t Traversal, eq Equals, u, z T, path []T) []T {
|
||||
if t.Visited(u) {
|
||||
return path
|
||||
}
|
||||
for _, v := range t.Edges(u) {
|
||||
if eq(v, z) {
|
||||
return append(path, z, u)
|
||||
}
|
||||
if p := dfsRecursive(t, eq, v, z, path); len(p) > 0 {
|
||||
return append(p, u)
|
||||
}
|
||||
}
|
||||
return path
|
||||
}
|
||||
+271
@@ -0,0 +1,271 @@
|
||||
// Copyright 2016 The OPA Authors. All rights reserved.
|
||||
// Use of this source code is governed by an Apache2
|
||||
// license that can be found in the LICENSE file.
|
||||
|
||||
package util
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// T is a concise way to refer to T.
|
||||
type T any
|
||||
|
||||
type Hasher interface {
|
||||
Hash() int
|
||||
}
|
||||
|
||||
type hashEntry[K any, V any] struct {
|
||||
k K
|
||||
v V
|
||||
next *hashEntry[K, V]
|
||||
}
|
||||
|
||||
// TypedHashMap represents a key/value map.
|
||||
type TypedHashMap[K any, V any] struct {
|
||||
keq func(K, K) bool
|
||||
veq func(V, V) bool
|
||||
khash func(K) int
|
||||
vhash func(V) int
|
||||
def V
|
||||
table map[int]*hashEntry[K, V]
|
||||
size int
|
||||
}
|
||||
|
||||
// NewTypedHashMap returns a new empty TypedHashMap.
|
||||
func NewTypedHashMap[K any, V any](keq func(K, K) bool, veq func(V, V) bool, khash func(K) int, vhash func(V) int, def V) *TypedHashMap[K, V] {
|
||||
return &TypedHashMap[K, V]{
|
||||
keq: keq,
|
||||
veq: veq,
|
||||
khash: khash,
|
||||
vhash: vhash,
|
||||
def: def,
|
||||
table: make(map[int]*hashEntry[K, V]),
|
||||
size: 0,
|
||||
}
|
||||
}
|
||||
|
||||
// HashMap represents a key/value map.
|
||||
type HashMap = TypedHashMap[T, T]
|
||||
|
||||
// NewHashMap returns a new empty HashMap.
|
||||
func NewHashMap(eq func(T, T) bool, hash func(T) int) *HashMap {
|
||||
return &HashMap{
|
||||
keq: eq,
|
||||
veq: eq,
|
||||
khash: hash,
|
||||
vhash: hash,
|
||||
def: nil,
|
||||
table: make(map[int]*hashEntry[T, T]),
|
||||
size: 0,
|
||||
}
|
||||
}
|
||||
|
||||
// Copy returns a shallow copy of this HashMap.
|
||||
func (h *TypedHashMap[K, V]) Copy() *TypedHashMap[K, V] {
|
||||
cpy := NewTypedHashMap(h.keq, h.veq, h.khash, h.vhash, h.def)
|
||||
h.Iter(func(k K, v V) bool {
|
||||
cpy.Put(k, v)
|
||||
return false
|
||||
})
|
||||
return cpy
|
||||
}
|
||||
|
||||
// Equal returns true if this HashMap equals the other HashMap.
|
||||
// Two hash maps are equal if they contain the same key/value pairs.
|
||||
func (h *TypedHashMap[K, V]) Equal(other *TypedHashMap[K, V]) bool {
|
||||
if h.Len() != other.Len() {
|
||||
return false
|
||||
}
|
||||
return !h.Iter(func(k K, v V) bool {
|
||||
ov, ok := other.Get(k)
|
||||
if !ok {
|
||||
return true
|
||||
}
|
||||
return !h.veq(v, ov)
|
||||
})
|
||||
}
|
||||
|
||||
// Get returns the value for k.
|
||||
func (h *TypedHashMap[K, V]) Get(k K) (V, bool) {
|
||||
hash := h.khash(k)
|
||||
for entry := h.table[hash]; entry != nil; entry = entry.next {
|
||||
if h.keq(entry.k, k) {
|
||||
return entry.v, true
|
||||
}
|
||||
}
|
||||
return h.def, false
|
||||
}
|
||||
|
||||
// Delete removes the key k.
|
||||
func (h *TypedHashMap[K, V]) Delete(k K) {
|
||||
hash := h.khash(k)
|
||||
var prev *hashEntry[K, V]
|
||||
for entry := h.table[hash]; entry != nil; entry = entry.next {
|
||||
if h.keq(entry.k, k) {
|
||||
if prev != nil {
|
||||
prev.next = entry.next
|
||||
} else {
|
||||
h.table[hash] = entry.next
|
||||
}
|
||||
h.size--
|
||||
return
|
||||
}
|
||||
prev = entry
|
||||
}
|
||||
}
|
||||
|
||||
// Hash returns the hash code for this hash map.
|
||||
func (h *TypedHashMap[K, V]) Hash() int {
|
||||
var hash int
|
||||
h.Iter(func(k K, v V) bool {
|
||||
hash += h.khash(k) + h.vhash(v)
|
||||
return false
|
||||
})
|
||||
return hash
|
||||
}
|
||||
|
||||
// Iter invokes the iter function for each element in the HashMap.
|
||||
// If the iter function returns true, iteration stops and the return value is true.
|
||||
// If the iter function never returns true, iteration proceeds through all elements
|
||||
// and the return value is false.
|
||||
func (h *TypedHashMap[K, V]) Iter(iter func(K, V) bool) bool {
|
||||
for _, entry := range h.table {
|
||||
for ; entry != nil; entry = entry.next {
|
||||
if iter(entry.k, entry.v) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// Len returns the current size of this HashMap.
|
||||
func (h *TypedHashMap[K, V]) Len() int {
|
||||
return h.size
|
||||
}
|
||||
|
||||
// Put inserts a key/value pair into this HashMap. If the key is already present, the existing
|
||||
// value is overwritten.
|
||||
func (h *TypedHashMap[K, V]) Put(k K, v V) {
|
||||
hash := h.khash(k)
|
||||
head := h.table[hash]
|
||||
for entry := head; entry != nil; entry = entry.next {
|
||||
if h.keq(entry.k, k) {
|
||||
entry.v = v
|
||||
return
|
||||
}
|
||||
}
|
||||
h.table[hash] = &hashEntry[K, V]{k: k, v: v, next: head}
|
||||
h.size++
|
||||
}
|
||||
|
||||
func (h *TypedHashMap[K, V]) String() string {
|
||||
var buf []string
|
||||
h.Iter(func(k K, v V) bool {
|
||||
buf = append(buf, fmt.Sprintf("%v: %v", k, v))
|
||||
return false
|
||||
})
|
||||
return "{" + strings.Join(buf, ", ") + "}"
|
||||
}
|
||||
|
||||
// Update returns a new HashMap with elements from the other HashMap put into this HashMap.
|
||||
// If the other HashMap contains elements with the same key as this HashMap, the value
|
||||
// from the other HashMap overwrites the value from this HashMap.
|
||||
func (h *TypedHashMap[K, V]) Update(other *TypedHashMap[K, V]) *TypedHashMap[K, V] {
|
||||
updated := h.Copy()
|
||||
other.Iter(func(k K, v V) bool {
|
||||
updated.Put(k, v)
|
||||
return false
|
||||
})
|
||||
return updated
|
||||
}
|
||||
|
||||
type hasherEntry[K Hasher, V any] struct {
|
||||
k K
|
||||
v V
|
||||
next *hasherEntry[K, V]
|
||||
}
|
||||
|
||||
// HasherMap represents a simpler version of TypedHashMap that uses Hasher's
|
||||
// for keys, and requires only an equality function for keys. Ideally we'd have
|
||||
// and Equal method for all key types too, and we could get rid of that requirement.
|
||||
type HasherMap[K Hasher, V any] struct {
|
||||
keq func(K, K) bool
|
||||
table map[int]*hasherEntry[K, V]
|
||||
size int
|
||||
}
|
||||
|
||||
// NewHasherMap returns a new empty HasherMap.
|
||||
func NewHasherMap[K Hasher, V any](keq func(K, K) bool) *HasherMap[K, V] {
|
||||
return &HasherMap[K, V]{
|
||||
keq: keq,
|
||||
table: make(map[int]*hasherEntry[K, V]),
|
||||
size: 0,
|
||||
}
|
||||
}
|
||||
|
||||
// Get returns the value for k.
|
||||
func (h *HasherMap[K, V]) Get(k K) (V, bool) {
|
||||
for entry := h.table[k.Hash()]; entry != nil; entry = entry.next {
|
||||
if h.keq(entry.k, k) {
|
||||
return entry.v, true
|
||||
}
|
||||
}
|
||||
var zero V
|
||||
return zero, false
|
||||
}
|
||||
|
||||
// Put inserts a key/value pair into this HashMap. If the key is already present, the existing
|
||||
// value is overwritten.
|
||||
func (h *HasherMap[K, V]) Put(k K, v V) {
|
||||
hash := k.Hash()
|
||||
head := h.table[hash]
|
||||
for entry := head; entry != nil; entry = entry.next {
|
||||
if h.keq(entry.k, k) {
|
||||
entry.v = v
|
||||
return
|
||||
}
|
||||
}
|
||||
h.table[hash] = &hasherEntry[K, V]{k: k, v: v, next: head}
|
||||
h.size++
|
||||
}
|
||||
|
||||
// Delete removes the key k.
|
||||
func (h *HasherMap[K, V]) Delete(k K) {
|
||||
hash := k.Hash()
|
||||
var prev *hasherEntry[K, V]
|
||||
for entry := h.table[hash]; entry != nil; entry = entry.next {
|
||||
if h.keq(entry.k, k) {
|
||||
if prev != nil {
|
||||
prev.next = entry.next
|
||||
} else {
|
||||
h.table[hash] = entry.next
|
||||
}
|
||||
h.size--
|
||||
return
|
||||
}
|
||||
prev = entry
|
||||
}
|
||||
}
|
||||
|
||||
// Iter invokes the iter function for each element in the HasherMap.
|
||||
// If the iter function returns true, iteration stops and the return value is true.
|
||||
// If the iter function never returns true, iteration proceeds through all elements
|
||||
// and the return value is false.
|
||||
func (h *HasherMap[K, V]) Iter(iter func(K, V) bool) bool {
|
||||
for _, entry := range h.table {
|
||||
for ; entry != nil; entry = entry.next {
|
||||
if iter(entry.k, entry.v) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// Len returns the current size of this HashMap.
|
||||
func (h *HasherMap[K, V]) Len() int {
|
||||
return h.size
|
||||
}
|
||||
+195
@@ -0,0 +1,195 @@
|
||||
// Copyright 2016 The OPA Authors. All rights reserved.
|
||||
// Use of this source code is governed by an Apache2
|
||||
// license that can be found in the LICENSE file.
|
||||
|
||||
package util
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"reflect"
|
||||
"strconv"
|
||||
|
||||
"sigs.k8s.io/yaml"
|
||||
|
||||
"github.com/open-policy-agent/opa/v1/loader/extension"
|
||||
)
|
||||
|
||||
// UnmarshalJSON parses the JSON encoded data and stores the result in the value
|
||||
// pointed to by x.
|
||||
//
|
||||
// This function is intended to be used in place of the standard [json.Marshal]
|
||||
// function when [json.Number] is required.
|
||||
func UnmarshalJSON(bs []byte, x any) error {
|
||||
return unmarshalJSON(bs, x, true)
|
||||
}
|
||||
|
||||
func unmarshalJSON(bs []byte, x any, ext bool) error {
|
||||
decoder := NewJSONDecoder(bytes.NewBuffer(bs))
|
||||
if err := decoder.Decode(x); err != nil {
|
||||
if handler := extension.FindExtension(".json"); handler != nil && ext {
|
||||
return handler(bs, x)
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
||||
// Since decoder.Decode validates only the first json structure in bytes,
|
||||
// check if decoder has more bytes to consume to validate whole input bytes.
|
||||
tok, err := decoder.Token()
|
||||
if tok != nil {
|
||||
return fmt.Errorf("error: invalid character '%s' after top-level value", tok)
|
||||
}
|
||||
if err != nil && err != io.EOF {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// NewJSONDecoder returns a new decoder that reads from r.
|
||||
//
|
||||
// This function is intended to be used in place of the standard [json.NewDecoder]
|
||||
// when [json.Number] is required.
|
||||
func NewJSONDecoder(r io.Reader) *json.Decoder {
|
||||
decoder := json.NewDecoder(r)
|
||||
decoder.UseNumber()
|
||||
return decoder
|
||||
}
|
||||
|
||||
// MustUnmarshalJSON parse the JSON encoded data and returns the result.
|
||||
//
|
||||
// If the data cannot be decoded, this function will panic. This function is for
|
||||
// test purposes.
|
||||
func MustUnmarshalJSON(bs []byte) any {
|
||||
var x any
|
||||
if err := UnmarshalJSON(bs, &x); err != nil {
|
||||
panic(err)
|
||||
}
|
||||
return x
|
||||
}
|
||||
|
||||
// MustMarshalJSON returns the JSON encoding of x
|
||||
//
|
||||
// If the data cannot be encoded, this function will panic. This function is for
|
||||
// test purposes.
|
||||
func MustMarshalJSON(x any) []byte {
|
||||
bs, err := json.Marshal(x)
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
return bs
|
||||
}
|
||||
|
||||
// RoundTrip encodes to JSON, and decodes the result again.
|
||||
//
|
||||
// Thereby, it is converting its argument to the representation expected by
|
||||
// rego.Input and inmem's Write operations. Works with both references and
|
||||
// values.
|
||||
func RoundTrip(x *any) error {
|
||||
// Avoid round-tripping types that won't change as a result of
|
||||
// marshalling/unmarshalling, as even for those values, round-tripping
|
||||
// comes with a significant cost.
|
||||
if x == nil || !NeedsRoundTrip(*x) {
|
||||
return nil
|
||||
}
|
||||
|
||||
// For number types, we can write the json.Number representation
|
||||
// directly into x without marshalling to bytes and back.
|
||||
a := *x
|
||||
switch v := a.(type) {
|
||||
case int:
|
||||
*x = json.Number(strconv.Itoa(v))
|
||||
return nil
|
||||
case int8:
|
||||
*x = json.Number(strconv.FormatInt(int64(v), 10))
|
||||
return nil
|
||||
case int16:
|
||||
*x = json.Number(strconv.FormatInt(int64(v), 10))
|
||||
return nil
|
||||
case int32:
|
||||
*x = json.Number(strconv.FormatInt(int64(v), 10))
|
||||
return nil
|
||||
case int64:
|
||||
*x = json.Number(strconv.FormatInt(v, 10))
|
||||
return nil
|
||||
case uint:
|
||||
*x = json.Number(strconv.FormatUint(uint64(v), 10))
|
||||
return nil
|
||||
case uint8:
|
||||
*x = json.Number(strconv.FormatUint(uint64(v), 10))
|
||||
return nil
|
||||
case uint16:
|
||||
*x = json.Number(strconv.FormatUint(uint64(v), 10))
|
||||
return nil
|
||||
case uint32:
|
||||
*x = json.Number(strconv.FormatUint(uint64(v), 10))
|
||||
return nil
|
||||
case uint64:
|
||||
*x = json.Number(strconv.FormatUint(v, 10))
|
||||
return nil
|
||||
case float32:
|
||||
*x = json.Number(strconv.FormatFloat(float64(v), 'f', -1, 32))
|
||||
return nil
|
||||
case float64:
|
||||
*x = json.Number(strconv.FormatFloat(v, 'f', -1, 64))
|
||||
return nil
|
||||
}
|
||||
|
||||
bs, err := json.Marshal(x)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return UnmarshalJSON(bs, x)
|
||||
}
|
||||
|
||||
// NeedsRoundTrip returns true if the value won't change as a result of
|
||||
// a marshalling/unmarshalling round-trip. Since [RoundTrip] itself calls
|
||||
// this you normally don't need to call this function directly, unless you
|
||||
// want to make decisions based on the round-tripability of a value without
|
||||
// actually doing the round-trip.
|
||||
func NeedsRoundTrip(x any) bool {
|
||||
switch x.(type) {
|
||||
case nil, bool, string, json.Number:
|
||||
return false
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
// Reference returns a pointer to its argument unless the argument already is
|
||||
// a pointer. If the argument is **t, or ***t, etc, it will return *t.
|
||||
//
|
||||
// Used for preparing Go types (including pointers to structs) into values to be
|
||||
// put through [RoundTrip].
|
||||
func Reference(x any) *any {
|
||||
var y any
|
||||
rv := reflect.ValueOf(x)
|
||||
if rv.Kind() == reflect.Pointer {
|
||||
return Reference(rv.Elem().Interface())
|
||||
}
|
||||
if rv.Kind() != reflect.Invalid {
|
||||
y = rv.Interface()
|
||||
return &y
|
||||
}
|
||||
return &x
|
||||
}
|
||||
|
||||
// Unmarshal decodes a YAML, JSON or JSON extension value into the specified type.
|
||||
func Unmarshal(bs []byte, v any) error {
|
||||
if len(bs) > 2 && bs[0] == 0xef && bs[1] == 0xbb && bs[2] == 0xbf {
|
||||
bs = bs[3:] // Strip UTF-8 BOM, see https://www.rfc-editor.org/rfc/rfc8259#section-8.1
|
||||
}
|
||||
|
||||
if json.Valid(bs) {
|
||||
return unmarshalJSON(bs, v, false)
|
||||
}
|
||||
nbs, err := yaml.YAMLToJSON(bs)
|
||||
if err == nil {
|
||||
return unmarshalJSON(nbs, v, false)
|
||||
}
|
||||
// not json or yaml: try extensions
|
||||
if handler := extension.FindExtension(".json"); handler != nil {
|
||||
return handler(bs, v)
|
||||
}
|
||||
return err
|
||||
}
|
||||
+34
@@ -0,0 +1,34 @@
|
||||
package util
|
||||
|
||||
import (
|
||||
"cmp"
|
||||
"slices"
|
||||
)
|
||||
|
||||
// Keys returns a slice of keys from any map.
|
||||
func Keys[M ~map[K]V, K comparable, V any](m M) []K {
|
||||
r := make([]K, 0, len(m))
|
||||
for k := range m {
|
||||
r = append(r, k)
|
||||
}
|
||||
return r
|
||||
}
|
||||
|
||||
// KeysSorted returns a slice of keys from any map, sorted in ascending order.
|
||||
func KeysSorted[M ~map[K]V, K cmp.Ordered, V any](m M) []K {
|
||||
r := make([]K, 0, len(m))
|
||||
for k := range m {
|
||||
r = append(r, k)
|
||||
}
|
||||
slices.Sort(r)
|
||||
return r
|
||||
}
|
||||
|
||||
// Values returns a slice of values from any map. Copied from golang.org/x/exp/maps.
|
||||
func Values[M ~map[K]V, K comparable, V any](m M) []V {
|
||||
r := make([]V, 0, len(m))
|
||||
for _, v := range m {
|
||||
r = append(r, v)
|
||||
}
|
||||
return r
|
||||
}
|
||||
+178
@@ -0,0 +1,178 @@
|
||||
package util
|
||||
|
||||
import (
|
||||
"slices"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"unsafe"
|
||||
)
|
||||
|
||||
// SyncPool is a generic sync.Pool for type T, providing some convenience
|
||||
// over sync.Pool directly: [SyncPool.Put] ensures that nil values are not
|
||||
// put into the pool, and [SyncPool.Get] returns a pointer to T without having
|
||||
// to do a type assertion at the call site.
|
||||
type SyncPool[T any] struct {
|
||||
pool sync.Pool
|
||||
}
|
||||
|
||||
func NewSyncPool[T any]() *SyncPool[T] {
|
||||
return &SyncPool[T]{
|
||||
pool: sync.Pool{
|
||||
New: func() any {
|
||||
return new(T)
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func (p *SyncPool[T]) Get() *T {
|
||||
return p.pool.Get().(*T)
|
||||
}
|
||||
|
||||
func (p *SyncPool[T]) Put(x *T) {
|
||||
if x != nil {
|
||||
p.pool.Put(x)
|
||||
}
|
||||
}
|
||||
|
||||
// NewPtrSlice returns a slice of pointers to T with length n,
|
||||
// with only 2 allocations performed no matter the size of n.
|
||||
// See:
|
||||
// https://gist.github.com/CAFxX/e96e8a5c3841d152f16d266a1fe7f8bd#slices-of-pointers
|
||||
func NewPtrSlice[T any](n int) []*T {
|
||||
return GrowPtrSlice[T](nil, n)
|
||||
}
|
||||
|
||||
// GrowPtrSlice appends n elements to the slice, each pointing to
|
||||
// a newly-allocated T. The resulting slice has length equal to len(s)+n.
|
||||
//
|
||||
// It performs at most 2 allocations, regardless of n.
|
||||
func GrowPtrSlice[T any](s []*T, n int) []*T {
|
||||
s = slices.Grow(s, n)
|
||||
p := make([]T, n)
|
||||
for i := range n {
|
||||
s = append(s, &p[i])
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
// Allocation free conversion from []byte to string (unsafe)
|
||||
// Note that the byte slice must not be modified after conversion
|
||||
func ByteSliceToString(bs []byte) string {
|
||||
return unsafe.String(unsafe.SliceData(bs), len(bs))
|
||||
}
|
||||
|
||||
// Allocation free conversion from ~string to []byte (unsafe)
|
||||
// Note that the byte slice must not be modified after conversion
|
||||
func StringToByteSlice[T ~string](s T) []byte {
|
||||
return unsafe.Slice(unsafe.StringData(string(s)), len(s))
|
||||
}
|
||||
|
||||
// NumDigitsInt returns the number of digits in n.
|
||||
// This is useful for pre-allocating buffers for string conversion.
|
||||
func NumDigitsInt(n int) int {
|
||||
return NumDigitsInt64(int64(n))
|
||||
}
|
||||
|
||||
// NumDigitsInt64 returns the number of digits in n.
|
||||
// This is useful for pre-allocating buffers for string conversion.
|
||||
func NumDigitsInt64(n int64) int {
|
||||
if n == 0 {
|
||||
return 1
|
||||
}
|
||||
|
||||
if n < 0 {
|
||||
n = -n
|
||||
}
|
||||
|
||||
count := 0
|
||||
for n > 0 {
|
||||
n /= 10
|
||||
count++
|
||||
}
|
||||
return count
|
||||
}
|
||||
|
||||
// NumDigitsUint returns the number of digits in n.
|
||||
// This is useful for pre-allocating buffers for string conversion.
|
||||
func NumDigitsUint(n uint64) int {
|
||||
if n == 0 {
|
||||
return 1
|
||||
}
|
||||
|
||||
count := 0
|
||||
for n > 0 {
|
||||
n /= 10
|
||||
count++
|
||||
}
|
||||
return count
|
||||
}
|
||||
|
||||
// AppendInt is a less messy version of strconv.AppendInt for base 10 ints.
|
||||
func AppendInt(buf []byte, n int) []byte {
|
||||
return strconv.AppendInt(buf, int64(n), 10)
|
||||
}
|
||||
|
||||
// SplitMap calls fn for each delim-separated part of text and returns a slice of the results.
|
||||
// Cheaper than calling fn on strings.Split(text, delim), as it avoids allocating an intermediate slice of strings.
|
||||
func SplitMap[T any](text string, delim string, fn func(string) T) []T {
|
||||
sl := make([]T, 0, strings.Count(text, delim)+1)
|
||||
for s := range strings.SplitSeq(text, delim) {
|
||||
sl = append(sl, fn(s))
|
||||
}
|
||||
return sl
|
||||
}
|
||||
|
||||
// SlicePool is a pool for (pointers to) slices of type T.
|
||||
// It uses sync.Pool to pool the slices, and grows them as needed.
|
||||
type SlicePool[T any] struct {
|
||||
pool sync.Pool
|
||||
}
|
||||
|
||||
// NewSlicePool creates a new SlicePool for slices of type T with the given initial length.
|
||||
// This number is only a hint, as the slices will grow as needed. For best performance, store
|
||||
// slices of similar lengths in the same pool.
|
||||
func NewSlicePool[T any](length int) *SlicePool[T] {
|
||||
return &SlicePool[T]{
|
||||
pool: sync.Pool{
|
||||
New: func() any {
|
||||
s := make([]T, length)
|
||||
return &s
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// Get returns a pointer to a slice of type T with the given length
|
||||
// from the pool. The slice capacity will grow as needed to accommodate
|
||||
// the requested length. The returned slice will have all its elements
|
||||
// set to the zero value of T. Returns a pointer to avoid allocating.
|
||||
func (sp *SlicePool[T]) Get(length int) *[]T {
|
||||
s := sp.pool.Get().(*[]T)
|
||||
d := *s
|
||||
|
||||
if cap(d) < length {
|
||||
d = slices.Grow(d, length)
|
||||
}
|
||||
|
||||
d = d[:length] // reslice to requested length, while keeping capacity
|
||||
|
||||
clear(d)
|
||||
|
||||
*s = d
|
||||
return s
|
||||
}
|
||||
|
||||
// Put returns a pointer to a slice of type T to the pool.
|
||||
func (sp *SlicePool[T]) Put(s *[]T) {
|
||||
if s != nil {
|
||||
sp.pool.Put(s)
|
||||
}
|
||||
}
|
||||
|
||||
// SortedFunc is simply a shorthand for [slices.SortFunc] which also returns the sorted slice.
|
||||
func SortedFunc[T any, S ~[]T](s S, cmp func(a, b T) int) S {
|
||||
slices.SortFunc(s, cmp)
|
||||
return s
|
||||
}
|
||||
+113
@@ -0,0 +1,113 @@
|
||||
// Copyright 2017 The OPA Authors. All rights reserved.
|
||||
// Use of this source code is governed by an Apache2
|
||||
// license that can be found in the LICENSE file.
|
||||
|
||||
package util
|
||||
|
||||
// LIFO represents a simple LIFO queue.
|
||||
type LIFO struct {
|
||||
top *queueNode
|
||||
size int
|
||||
}
|
||||
|
||||
type queueNode struct {
|
||||
v T
|
||||
next *queueNode
|
||||
}
|
||||
|
||||
// NewLIFO returns a new LIFO queue containing elements ts starting with the
|
||||
// left-most argument at the bottom.
|
||||
func NewLIFO(ts ...T) *LIFO {
|
||||
s := &LIFO{}
|
||||
for i := range ts {
|
||||
s.Push(ts[i])
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
// Push adds a new element onto the LIFO.
|
||||
func (s *LIFO) Push(t T) {
|
||||
node := &queueNode{v: t, next: s.top}
|
||||
s.top = node
|
||||
s.size++
|
||||
}
|
||||
|
||||
// Peek returns the top of the LIFO. If LIFO is empty, returns nil, false.
|
||||
func (s *LIFO) Peek() (T, bool) {
|
||||
if s.top == nil {
|
||||
return nil, false
|
||||
}
|
||||
return s.top.v, true
|
||||
}
|
||||
|
||||
// Pop returns the top of the LIFO and removes it. If LIFO is empty returns
|
||||
// nil, false.
|
||||
func (s *LIFO) Pop() (T, bool) {
|
||||
if s.top == nil {
|
||||
return nil, false
|
||||
}
|
||||
node := s.top
|
||||
s.top = node.next
|
||||
s.size--
|
||||
return node.v, true
|
||||
}
|
||||
|
||||
// Size returns the size of the LIFO.
|
||||
func (s *LIFO) Size() int {
|
||||
return s.size
|
||||
}
|
||||
|
||||
// FIFO represents a simple FIFO queue.
|
||||
type FIFO struct {
|
||||
front *queueNode
|
||||
back *queueNode
|
||||
size int
|
||||
}
|
||||
|
||||
// NewFIFO returns a new FIFO queue containing elements ts starting with the
|
||||
// left-most argument at the front.
|
||||
func NewFIFO(ts ...T) *FIFO {
|
||||
s := &FIFO{}
|
||||
for i := range ts {
|
||||
s.Push(ts[i])
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
// Push adds a new element onto the LIFO.
|
||||
func (s *FIFO) Push(t T) {
|
||||
node := &queueNode{v: t, next: nil}
|
||||
if s.front == nil {
|
||||
s.front = node
|
||||
s.back = node
|
||||
} else {
|
||||
s.back.next = node
|
||||
s.back = node
|
||||
}
|
||||
s.size++
|
||||
}
|
||||
|
||||
// Peek returns the top of the LIFO. If LIFO is empty, returns nil, false.
|
||||
func (s *FIFO) Peek() (T, bool) {
|
||||
if s.front == nil {
|
||||
return nil, false
|
||||
}
|
||||
return s.front.v, true
|
||||
}
|
||||
|
||||
// Pop returns the top of the LIFO and removes it. If LIFO is empty returns
|
||||
// nil, false.
|
||||
func (s *FIFO) Pop() (T, bool) {
|
||||
if s.front == nil {
|
||||
return nil, false
|
||||
}
|
||||
node := s.front
|
||||
s.front = node.next
|
||||
s.size--
|
||||
return node.v, true
|
||||
}
|
||||
|
||||
// Size returns the size of the LIFO.
|
||||
func (s *FIFO) Size() int {
|
||||
return s.size
|
||||
}
|
||||
+60
@@ -0,0 +1,60 @@
|
||||
package util
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"compress/gzip"
|
||||
"errors"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"github.com/open-policy-agent/opa/v1/util/decoding"
|
||||
)
|
||||
|
||||
var gzipReaderPool = NewSyncPool[gzip.Reader]()
|
||||
|
||||
// Note(philipc): Originally taken from server/server.go
|
||||
// The DecodingLimitHandler handles setting the max size limits in the context.
|
||||
// This function enforces those limits. For gzip payloads, we use a LimitReader
|
||||
// to ensure we don't decompress more than the allowed maximum, preventing
|
||||
// memory exhaustion from forged gzip trailers.
|
||||
func ReadMaybeCompressedBody(r *http.Request) ([]byte, error) {
|
||||
length := r.ContentLength
|
||||
if maxLenConf, ok := decoding.GetServerDecodingMaxLen(r.Context()); ok {
|
||||
length = maxLenConf
|
||||
}
|
||||
|
||||
content, err := io.ReadAll(io.LimitReader(r.Body, length))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if strings.Contains(r.Header.Get("Content-Encoding"), "gzip") {
|
||||
gzipMaxLength, _ := decoding.GetServerDecodingGzipMaxLen(r.Context())
|
||||
|
||||
gzReader := gzipReaderPool.Get()
|
||||
defer func() {
|
||||
gzReader.Close()
|
||||
gzipReaderPool.Put(gzReader)
|
||||
}()
|
||||
|
||||
if err := gzReader.Reset(bytes.NewReader(content)); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
decompressed := bytes.NewBuffer(make([]byte, 0, len(content)))
|
||||
limitReader := io.LimitReader(gzReader, gzipMaxLength+1)
|
||||
if _, err := decompressed.ReadFrom(limitReader); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if int64(decompressed.Len()) > gzipMaxLength {
|
||||
return nil, errors.New("gzip payload too large")
|
||||
}
|
||||
|
||||
return decompressed.Bytes(), nil
|
||||
}
|
||||
|
||||
// Request was not compressed; return the content bytes.
|
||||
return content, nil
|
||||
}
|
||||
+13
@@ -0,0 +1,13 @@
|
||||
package util
|
||||
|
||||
import "strings"
|
||||
|
||||
// WithPrefix ensures that the string s starts with the given prefix.
|
||||
// If s already starts with prefix, it is returned unchanged.
|
||||
func WithPrefix(s, prefix string) string {
|
||||
if strings.HasPrefix(s, prefix) {
|
||||
return s
|
||||
}
|
||||
|
||||
return prefix + s
|
||||
}
|
||||
+48
@@ -0,0 +1,48 @@
|
||||
package util
|
||||
|
||||
import "time"
|
||||
|
||||
// TimerWithCancel exists because of memory leaks when using
|
||||
// time.After in select statements. Instead, we now manually create timers,
|
||||
// wait on them, and manually free them.
|
||||
//
|
||||
// See this for more details:
|
||||
// https://www.arangodb.com/2020/09/a-story-of-a-memory-leak-in-go-how-to-properly-use-time-after/
|
||||
//
|
||||
// Note: This issue is fixed in Go 1.23, but this fix helps us until then.
|
||||
//
|
||||
// Warning: the cancel cannot be done concurrent to reading, everything should
|
||||
// work in the same goroutine.
|
||||
//
|
||||
// Example:
|
||||
//
|
||||
// for retries := 0; true; retries++ {
|
||||
//
|
||||
// ...main logic...
|
||||
//
|
||||
// timer, cancel := utils.TimerWithCancel(utils.Backoff(retries))
|
||||
// select {
|
||||
// case <-ctx.Done():
|
||||
// cancel()
|
||||
// return ctx.Err()
|
||||
// case <-timer.C:
|
||||
// continue
|
||||
// }
|
||||
// }
|
||||
func TimerWithCancel(delay time.Duration) (*time.Timer, func()) {
|
||||
timer := time.NewTimer(delay)
|
||||
|
||||
return timer, func() {
|
||||
// Note: The Stop function returns:
|
||||
// - true: if the timer is active. (no draining required)
|
||||
// - false: if the timer was already stopped or fired/expired.
|
||||
// In this case the channel should be drained to prevent memory
|
||||
// leaks only if it is not empty.
|
||||
// This operation is safe only if the cancel function is
|
||||
// used in same goroutine. Concurrent reading or canceling may
|
||||
// cause deadlock.
|
||||
if !timer.Stop() && len(timer.C) > 0 {
|
||||
<-timer.C
|
||||
}
|
||||
}
|
||||
}
|
||||
+34
@@ -0,0 +1,34 @@
|
||||
// Copyright 2020 The OPA Authors. All rights reserved.
|
||||
// Use of this source code is governed by an Apache2
|
||||
// license that can be found in the LICENSE file.
|
||||
|
||||
package util
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"time"
|
||||
)
|
||||
|
||||
// WaitFunc will call passed function at an interval and return nil
|
||||
// as soon this function returns true.
|
||||
// If timeout is reached before the passed in function returns true
|
||||
// an error is returned.
|
||||
func WaitFunc(fun func() bool, interval, timeout time.Duration) error {
|
||||
if fun() {
|
||||
return nil
|
||||
}
|
||||
ticker := time.NewTicker(interval)
|
||||
timer := time.NewTimer(timeout)
|
||||
defer ticker.Stop()
|
||||
defer timer.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-timer.C:
|
||||
return errors.New("timeout")
|
||||
case <-ticker.C:
|
||||
if fun() {
|
||||
return nil
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user