|
| 1 | +package sandbox |
| 2 | + |
| 3 | +import ( |
| 4 | + "sync" |
| 5 | + |
| 6 | + "github.com/e2b-dev/infra/packages/shared/pkg/smap" |
| 7 | +) |
| 8 | + |
| 9 | +type MapSubscriber interface { |
| 10 | + OnInsert(sandbox *Sandbox) |
| 11 | + OnRemove(sandboxID string) |
| 12 | +} |
| 13 | + |
| 14 | +type Map struct { |
| 15 | + sandboxes *smap.Map[*Sandbox] |
| 16 | + |
| 17 | + subs []MapSubscriber |
| 18 | + subsLock sync.RWMutex |
| 19 | +} |
| 20 | + |
| 21 | +func (m *Map) Subscribe(subscriber MapSubscriber) { |
| 22 | + m.subsLock.Lock() |
| 23 | + defer m.subsLock.Unlock() |
| 24 | + |
| 25 | + m.subs = append(m.subs, subscriber) |
| 26 | +} |
| 27 | + |
| 28 | +func (m *Map) trigger(fn func(MapSubscriber)) { |
| 29 | + m.subsLock.RLock() |
| 30 | + defer m.subsLock.RUnlock() |
| 31 | + |
| 32 | + for _, subscriber := range m.subs { |
| 33 | + fn(subscriber) |
| 34 | + } |
| 35 | +} |
| 36 | + |
| 37 | +func (m *Map) Items() map[string]*Sandbox { |
| 38 | + return m.sandboxes.Items() |
| 39 | +} |
| 40 | + |
| 41 | +func (m *Map) Count() int { |
| 42 | + return m.sandboxes.Count() |
| 43 | +} |
| 44 | + |
| 45 | +func (m *Map) Get(sandboxID string) (*Sandbox, bool) { |
| 46 | + return m.sandboxes.Get(sandboxID) |
| 47 | +} |
| 48 | + |
| 49 | +func (m *Map) Insert(sbx *Sandbox) { |
| 50 | + m.sandboxes.Insert(sbx.Runtime.SandboxID, sbx) |
| 51 | + |
| 52 | + go m.trigger(func(s MapSubscriber) { |
| 53 | + s.OnInsert(sbx) |
| 54 | + }) |
| 55 | +} |
| 56 | + |
| 57 | +func (m *Map) Remove(sandboxID string) { |
| 58 | + m.sandboxes.Remove(sandboxID) |
| 59 | + |
| 60 | + go m.trigger(func(s MapSubscriber) { |
| 61 | + s.OnRemove(sandboxID) |
| 62 | + }) |
| 63 | +} |
| 64 | + |
| 65 | +func (m *Map) RemoveByExecutionID(sandboxID, executionID string) { |
| 66 | + removed := m.sandboxes.RemoveCb(sandboxID, func(_ string, v *Sandbox, exists bool) bool { |
| 67 | + if !exists { |
| 68 | + return false |
| 69 | + } |
| 70 | + |
| 71 | + if v == nil { |
| 72 | + return false |
| 73 | + } |
| 74 | + |
| 75 | + return v.Runtime.ExecutionID == executionID |
| 76 | + }) |
| 77 | + |
| 78 | + if removed { |
| 79 | + go m.trigger(func(s MapSubscriber) { |
| 80 | + s.OnRemove(sandboxID) |
| 81 | + }) |
| 82 | + } |
| 83 | +} |
| 84 | + |
| 85 | +func NewSandboxesMap() *Map { |
| 86 | + return &Map{sandboxes: smap.New[*Sandbox]()} |
| 87 | +} |
0 commit comments