internal/xsync/map.go - master - githubmirror/ydb-go-sdk

internal/xsync/map.go

78 строк1 KB

Aleksey Myasnikov

refactor(conn): Get/Put ref-counting with balancer quarantine (#2235)

15 июл 2026, 17:24

15 июл 2026, 17:246dfb4a0

package xsync

import (
	"fmt"
	"sync"
)

type Map[K comparable, V any] struct {
	m sync.Map
}

func (m *Map[K, V]) Get(key K) (value V, ok bool) {
	v, ok := m.m.Load(key)
	if !ok {
		return value, false
	}
	value, ok = v.(V)

	return value, ok
}

func (m *Map[K, V]) Must(key K) (value V) {
	v, ok := m.m.Load(key)
	if !ok {
		panic(fmt.Sprintf("unexpected key = %v", key))
	}
	value, ok = v.(V)
	if !ok {
		panic(fmt.Sprintf("unexpected type of value = %T", v))
	}

	return value
}

func (m *Map[K, V]) Has(key K) bool {
	_, ok := m.m.Load(key)

	return ok
}

func (m *Map[K, V]) Set(key K, value V) {
	m.m.Store(key, value)
}

func (m *Map[K, V]) Delete(key K) bool {
	_, exists := m.m.LoadAndDelete(key)

	return exists
}

func (m *Map[K, V]) Extract(key K) (value V, ok bool) {
	v, exists := m.m.LoadAndDelete(key)
	if !exists {
		return value, false
	}

	value, ok = v.(V)

	return value, ok
}

func (m *Map[K, V]) Range(f func(key K, value V) bool) {
	m.m.Range(func(k, v any) bool {
		return f(k.(K), v.(V)) //nolint:forcetypeassert
	})
}

func (m *Map[K, V]) Clear() (removed int) {
	m.m.Range(func(k, v any) bool {
		removed++

		m.m.Delete(k)

		return true
	})

	return removed
}