mirror of
https://github.com/caddyserver/caddy.git
synced 2024-12-25 05:03:47 +03:00
227 lines
6.6 KiB
Go
227 lines
6.6 KiB
Go
// Copyright 2015 Matthew Holt and The Caddy Authors
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
// you may not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
|
|
package caddy
|
|
|
|
import (
|
|
"fmt"
|
|
"sync"
|
|
"sync/atomic"
|
|
)
|
|
|
|
// UsagePool is a thread-safe map that pools values
|
|
// based on usage (reference counting). Values are
|
|
// only inserted if they do not already exist. There
|
|
// are two ways to add values to the pool:
|
|
//
|
|
// 1. LoadOrStore will increment usage and store the
|
|
// value immediately if it does not already exist.
|
|
// 2. LoadOrNew will atomically check for existence
|
|
// and construct the value immediately if it does
|
|
// not already exist, or increment the usage
|
|
// otherwise, then store that value in the pool.
|
|
// When the constructed value is finally deleted
|
|
// from the pool (when its usage reaches 0), it
|
|
// will be cleaned up by calling Destruct().
|
|
//
|
|
// The use of LoadOrNew allows values to be created
|
|
// and reused and finally cleaned up only once, even
|
|
// though they may have many references throughout
|
|
// their lifespan. This is helpful, for example, when
|
|
// sharing thread-safe io.Writers that you only want
|
|
// to open and close once.
|
|
//
|
|
// There is no way to overwrite existing keys in the
|
|
// pool without first deleting it as many times as it
|
|
// was stored. Deleting too many times will panic.
|
|
//
|
|
// The implementation does not use a sync.Pool because
|
|
// UsagePool needs additional atomicity to run the
|
|
// constructor functions when creating a new value when
|
|
// LoadOrNew is used. (We could probably use sync.Pool
|
|
// but we'd still have to layer our own additional locks
|
|
// on top.)
|
|
//
|
|
// An empty UsagePool is NOT safe to use; always call
|
|
// NewUsagePool() to make a new one.
|
|
type UsagePool struct {
|
|
sync.RWMutex
|
|
pool map[any]*usagePoolVal
|
|
}
|
|
|
|
// NewUsagePool returns a new usage pool that is ready to use.
|
|
func NewUsagePool() *UsagePool {
|
|
return &UsagePool{
|
|
pool: make(map[any]*usagePoolVal),
|
|
}
|
|
}
|
|
|
|
// LoadOrNew loads the value associated with key from the pool if it
|
|
// already exists. If the key doesn't exist, it will call construct
|
|
// to create a new value and then stores that in the pool. An error
|
|
// is only returned if the constructor returns an error. The loaded
|
|
// or constructed value is returned. The loaded return value is true
|
|
// if the value already existed and was loaded, or false if it was
|
|
// newly constructed.
|
|
func (up *UsagePool) LoadOrNew(key any, construct Constructor) (value any, loaded bool, err error) {
|
|
var upv *usagePoolVal
|
|
up.Lock()
|
|
upv, loaded = up.pool[key]
|
|
if loaded {
|
|
atomic.AddInt32(&upv.refs, 1)
|
|
up.Unlock()
|
|
upv.RLock()
|
|
value = upv.value
|
|
err = upv.err
|
|
upv.RUnlock()
|
|
} else {
|
|
upv = &usagePoolVal{refs: 1}
|
|
upv.Lock()
|
|
up.pool[key] = upv
|
|
up.Unlock()
|
|
value, err = construct()
|
|
if err == nil {
|
|
upv.value = value
|
|
} else {
|
|
upv.err = err
|
|
up.Lock()
|
|
// this *should* be safe, I think, because we have a
|
|
// write lock on upv, but we might also need to ensure
|
|
// that upv.err is nil before doing this, since we
|
|
// released the write lock on up during construct...
|
|
// but then again it's also after midnight...
|
|
delete(up.pool, key)
|
|
up.Unlock()
|
|
}
|
|
upv.Unlock()
|
|
}
|
|
return
|
|
}
|
|
|
|
// LoadOrStore loads the value associated with key from the pool if it
|
|
// already exists, or stores it if it does not exist. It returns the
|
|
// value that was either loaded or stored, and true if the value already
|
|
// existed and was loaded, false if the value didn't exist and was stored.
|
|
func (up *UsagePool) LoadOrStore(key, val any) (value any, loaded bool) {
|
|
var upv *usagePoolVal
|
|
up.Lock()
|
|
upv, loaded = up.pool[key]
|
|
if loaded {
|
|
atomic.AddInt32(&upv.refs, 1)
|
|
up.Unlock()
|
|
upv.Lock()
|
|
if upv.err == nil {
|
|
value = upv.value
|
|
} else {
|
|
upv.value = val
|
|
upv.err = nil
|
|
}
|
|
upv.Unlock()
|
|
} else {
|
|
upv = &usagePoolVal{refs: 1, value: val}
|
|
up.pool[key] = upv
|
|
up.Unlock()
|
|
value = val
|
|
}
|
|
return
|
|
}
|
|
|
|
// Range iterates the pool similarly to how sync.Map.Range() does:
|
|
// it calls f for every key in the pool, and if f returns false,
|
|
// iteration is stopped. Ranging does not affect usage counts.
|
|
//
|
|
// This method is somewhat naive and acquires a read lock on the
|
|
// entire pool during iteration, so do your best to make f() really
|
|
// fast, m'kay?
|
|
func (up *UsagePool) Range(f func(key, value any) bool) {
|
|
up.RLock()
|
|
defer up.RUnlock()
|
|
for key, upv := range up.pool {
|
|
upv.RLock()
|
|
if upv.err != nil {
|
|
upv.RUnlock()
|
|
continue
|
|
}
|
|
val := upv.value
|
|
upv.RUnlock()
|
|
if !f(key, val) {
|
|
break
|
|
}
|
|
}
|
|
}
|
|
|
|
// Delete decrements the usage count for key and removes the
|
|
// value from the underlying map if the usage is 0. It returns
|
|
// true if the usage count reached 0 and the value was deleted.
|
|
// It panics if the usage count drops below 0; always call
|
|
// Delete precisely as many times as LoadOrStore.
|
|
func (up *UsagePool) Delete(key any) (deleted bool, err error) {
|
|
up.Lock()
|
|
upv, ok := up.pool[key]
|
|
if !ok {
|
|
up.Unlock()
|
|
return false, nil
|
|
}
|
|
refs := atomic.AddInt32(&upv.refs, -1)
|
|
if refs == 0 {
|
|
delete(up.pool, key)
|
|
up.Unlock()
|
|
upv.RLock()
|
|
val := upv.value
|
|
upv.RUnlock()
|
|
if destructor, ok := val.(Destructor); ok {
|
|
err = destructor.Destruct()
|
|
}
|
|
deleted = true
|
|
} else {
|
|
up.Unlock()
|
|
if refs < 0 {
|
|
panic(fmt.Sprintf("deleted more than stored: %#v (usage: %d)",
|
|
upv.value, upv.refs))
|
|
}
|
|
}
|
|
return
|
|
}
|
|
|
|
// References returns the number of references (count of usages) to a
|
|
// key in the pool, and true if the key exists, or false otherwise.
|
|
func (up *UsagePool) References(key any) (int, bool) {
|
|
up.RLock()
|
|
upv, loaded := up.pool[key]
|
|
up.RUnlock()
|
|
if loaded {
|
|
// I wonder if it'd be safer to read this value during
|
|
// our lock on the UsagePool... guess we'll see...
|
|
refs := atomic.LoadInt32(&upv.refs)
|
|
return int(refs), true
|
|
}
|
|
return 0, false
|
|
}
|
|
|
|
// Constructor is a function that returns a new value
|
|
// that can destruct itself when it is no longer needed.
|
|
type Constructor func() (Destructor, error)
|
|
|
|
// Destructor is a value that can clean itself up when
|
|
// it is deallocated.
|
|
type Destructor interface {
|
|
Destruct() error
|
|
}
|
|
|
|
type usagePoolVal struct {
|
|
refs int32 // accessed atomically; must be 64-bit aligned for 32-bit systems
|
|
value any
|
|
err error
|
|
sync.RWMutex
|
|
}
|