mirror of https://github.com/hashicorp/consul
60 lines
1.2 KiB
Go
60 lines
1.2 KiB
Go
|
package proto
|
||
|
|
|
||
|
|
import (
|
||
|
|
"github.com/inconshreveable/muxado/proto/frame"
|
||
|
|
"sync"
|
||
|
|
)
|
||
|
|
|
||
|
|
const (
|
||
|
|
initMapCapacity = 128 // not too much extra memory wasted to avoid allocations
|
||
|
|
)
|
||
|
|
|
||
|
|
type StreamMap interface {
|
||
|
|
Get(frame.StreamId) (stream, bool)
|
||
|
|
Set(frame.StreamId, stream)
|
||
|
|
Delete(frame.StreamId)
|
||
|
|
Each(func(frame.StreamId, stream))
|
||
|
|
}
|
||
|
|
|
||
|
|
// ConcurrentStreamMap is a map of stream ids -> streams guarded by a read/write lock
|
||
|
|
type ConcurrentStreamMap struct {
|
||
|
|
sync.RWMutex
|
||
|
|
table map[frame.StreamId]stream
|
||
|
|
}
|
||
|
|
|
||
|
|
func (m *ConcurrentStreamMap) Get(id frame.StreamId) (s stream, ok bool) {
|
||
|
|
m.RLock()
|
||
|
|
s, ok = m.table[id]
|
||
|
|
m.RUnlock()
|
||
|
|
return
|
||
|
|
}
|
||
|
|
|
||
|
|
func (m *ConcurrentStreamMap) Set(id frame.StreamId, str stream) {
|
||
|
|
m.Lock()
|
||
|
|
m.table[id] = str
|
||
|
|
m.Unlock()
|
||
|
|
}
|
||
|
|
|
||
|
|
func (m *ConcurrentStreamMap) Delete(id frame.StreamId) {
|
||
|
|
m.Lock()
|
||
|
|
delete(m.table, id)
|
||
|
|
m.Unlock()
|
||
|
|
}
|
||
|
|
|
||
|
|
func (m *ConcurrentStreamMap) Each(fn func(frame.StreamId, stream)) {
|
||
|
|
m.Lock()
|
||
|
|
streams := make(map[frame.StreamId]stream, len(m.table))
|
||
|
|
for k, v := range m.table {
|
||
|
|
streams[k] = v
|
||
|
|
}
|
||
|
|
m.Unlock()
|
||
|
|
|
||
|
|
for id, str := range streams {
|
||
|
|
fn(id, str)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func NewConcurrentStreamMap() *ConcurrentStreamMap {
|
||
|
|
return &ConcurrentStreamMap{table: make(map[frame.StreamId]stream, initMapCapacity)}
|
||
|
|
}
|