diff --git a/mutex_map.go b/mutex_map.go index 4a91d345..26284746 100644 --- a/mutex_map.go +++ b/mutex_map.go @@ -34,6 +34,17 @@ func (m *MutexMap[K, V]) Delete(key K) { delete(m.real, key) } +// RunAndDelete runs a callback and removes the key afterwards +func (m *MutexMap[K, V]) RunAndDelete(key K, callback func(key K, value V)) { + m.Lock() + defer m.Unlock() + + if value, ok := m.real[key]; ok { + callback(key, value) + delete(m.real, key) + } +} + // Size returns the length of the internal map func (m *MutexMap[K, V]) Size() int { m.RLock() diff --git a/packet_resend_manager.go b/packet_resend_manager.go index 47b7028c..efa89d0a 100644 --- a/packet_resend_manager.go +++ b/packet_resend_manager.go @@ -90,10 +90,9 @@ func (p *PacketResendManager) Add(packet PacketInterface) { // Remove removes a packet from pool and stops it's timer func (p *PacketResendManager) Remove(sequenceID uint16) { - if cached, ok := p.pending.Get(sequenceID); ok { - cached.StopTimeoutTimer() - p.pending.Delete(sequenceID) - } + p.pending.RunAndDelete(sequenceID, func(key uint16, value *PendingPacket){ + value.StopTimeoutTimer() + }) } // Clear removes all packets from pool and stops their timers