rename the sessionMap to packetHandlerMap

packet_handler_map.go

@@ -0,0 +1,74 @@
+package quic
+
+import (
+	"sync"
+	"time"
+
+	"github.com/lucas-clemente/quic-go/internal/protocol"
+)
+
+type packetHandlerMap struct {
+	mutex sync.RWMutex
+
+	handlers map[string] /* string(ConnectionID)*/ packetHandler
+	closed   bool
+
+	deleteClosedSessionsAfter time.Duration
+}
+
+var _ packetHandlerManager = &packetHandlerMap{}
+
+func newPacketHandlerMap() packetHandlerManager {
+	return &packetHandlerMap{
+		handlers:                  make(map[string]packetHandler),
+		deleteClosedSessionsAfter: protocol.ClosedSessionDeleteTimeout,
+	}
+}
+
+func (h *packetHandlerMap) Get(id protocol.ConnectionID) (packetHandler, bool) {
+	h.mutex.RLock()
+	sess, ok := h.handlers[string(id)]
+	h.mutex.RUnlock()
+	return sess, ok
+}
+
+func (h *packetHandlerMap) Add(id protocol.ConnectionID, handler packetHandler) {
+	h.mutex.Lock()
+	h.handlers[string(id)] = handler
+	h.mutex.Unlock()
+}
+
+func (h *packetHandlerMap) Remove(id protocol.ConnectionID) {
+	h.mutex.Lock()
+	h.handlers[string(id)] = nil
+	h.mutex.Unlock()
+
+	time.AfterFunc(h.deleteClosedSessionsAfter, func() {
+		h.mutex.Lock()
+		delete(h.handlers, string(id))
+		h.mutex.Unlock()
+	})
+}
+
+func (h *packetHandlerMap) Close() {
+	h.mutex.Lock()
+	if h.closed {
+		h.mutex.Unlock()
+		return
+	}
+	h.closed = true
+
+	var wg sync.WaitGroup
+	for _, handler := range h.handlers {
+		if handler != nil {
+			wg.Add(1)
+			go func(handler packetHandler) {
+				// session.Close() blocks until the CONNECTION_CLOSE has been sent and the run-loop has stopped
+				_ = handler.Close(nil)
+				wg.Done()
+			}(handler)
+		}
+	}
+	h.mutex.Unlock()
+	wg.Wait()
+}