add GetAndSet to bitmap (#942)

* add GetAndSet to bitmap

* cleanup

* benchmark GetAndSet

Author: paulwe

logger/slog.go

@@ -1,3 +1,17 @@
+// Copyright 2023 LiveKit, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
 package logger
 
 import (

logger/zaputil/deferrer.go

@@ -17,6 +17,7 @@ package zaputil
 import (
 	"sync"
 
+	"go.uber.org/atomic"
 	"go.uber.org/zap"
 	"go.uber.org/zap/zapcore"
 )
@@ -70,25 +71,27 @@ type DeferredFieldResolver func(args ...any)
 
 func NewDeferrer() (*Deferrer, DeferredFieldResolver) {
 	buf := &Deferrer{}
-	var resolveOnce sync.Once
+	var resolved atomic.Bool
 	resolve := func(args ...any) {
-		resolveOnce.Do(func() {
-			fields := make([]zapcore.Field, 0, len(args))
-			for i := 0; i < len(args); i++ {
-				switch arg := args[i].(type) {
-				case zapcore.Field:
-					fields = append(fields, arg)
-				case string:
-					if i < len(args)-1 {
-						fields = append(fields, zap.Any(arg, args[i+1]))
-						i++
-					}
+		if resolved.Swap(true) {
+			return
+		}
+
+		fields := make([]zapcore.Field, 0, len(args))
+		for i := 0; i < len(args); i++ {
+			switch arg := args[i].(type) {
+			case zapcore.Field:
+				fields = append(fields, arg)
+			case string:
+				if i < len(args)-1 {
+					fields = append(fields, zap.Any(arg, args[i+1]))
+					i++
 				}
 			}
+		}
 
-			buf.fields = fields
-			buf.flush()
-		})
+		buf.fields = fields
+		buf.flush()
 	}
 	return buf, resolve
 }
@@ -99,8 +102,6 @@ type deferredValueCore struct {
 }
 
 func NewDeferredValueCore(core zapcore.Core, def *Deferrer) zapcore.Core {
-	def.mu.Lock()
-	defer def.mu.Unlock()
 	return &deferredValueCore{core, def}
 }
 

utils/bitmap.go

@@ -22,14 +22,11 @@ type bitmapNumber interface {
 
 type Bitmap[T bitmapNumber] struct {
 	bits []uint64
-	mask int
 }
 
 func NewBitmap[T bitmapNumber](size int) *Bitmap[T] {
-	numElems := 1 << bits.Len64(uint64(size+63)/64)
 	return &Bitmap[T]{
-		bits: make([]uint64, numElems),
-		mask: numElems - 1,
+		bits: make([]uint64, 1<<bits.Len64(uint64(size+63)/64)),
 	}
 }
 
@@ -38,61 +35,71 @@ func (b *Bitmap[T]) Len() int {
 }
 
 func (b *Bitmap[T]) Set(val T) {
-	s, o := b.getSlotAndOffset(val)
-	b.bits[s&b.mask] |= 1 << o
+	sm, s, o := b.getSlotAndOffset(val)
+	b.bits[s&sm] |= 1 << o
+}
+
+func (b *Bitmap[T]) GetAndSet(val T) bool {
+	sm, s, o := b.getSlotAndOffset(val)
+	prev := b.bits[s&sm]&(1<<o) != 0
+	b.bits[s&sm] |= 1 << o
+	return prev
 }
 
 func (b *Bitmap[T]) SetRange(min, max T) {
 	if max < min {
 		return
 	}
 
-	ls, rs, lo, ro := b.getSlotsAndOffsets(min, max)
+	sm, ls, rs, lo, ro := b.getSlotsAndOffsets(min, max)
 	if ls == rs {
-		b.bits[ls&b.mask] |= (((1 << (ro - lo + 1)) - 1) << lo)
+		b.bits[ls&sm] |= (((1 << (ro - lo + 1)) - 1) << lo)
 	} else {
-		b.bits[ls&b.mask] |= ^((1 << lo) - 1)
+		b.bits[ls&sm] |= ^((1 << lo) - 1)
 		for i := ls + 1; i < rs; i++ {
-			b.bits[i&b.mask] = ^uint64(0)
+			b.bits[i&sm] = ^uint64(0)
 		}
-		b.bits[rs&b.mask] |= (1 << (ro + 1)) - 1
+		b.bits[rs&sm] |= (1 << (ro + 1)) - 1
 	}
 }
 
 func (b *Bitmap[T]) Clear(val T) {
-	s, o := b.getSlotAndOffset(val)
-	b.bits[s&b.mask] &= ^(1 << o)
+	sm, s, o := b.getSlotAndOffset(val)
+	b.bits[s&sm] &= ^(1 << o)
 }
 
 func (b *Bitmap[T]) ClearRange(min, max T) {
 	if max < min {
 		return
 	}
 
-	ls, rs, lo, ro := b.getSlotsAndOffsets(min, max)
+	sm, ls, rs, lo, ro := b.getSlotsAndOffsets(min, max)
 	if ls == rs {
-		b.bits[ls&b.mask] &= ^(((1 << (ro - lo + 1)) - 1) << lo)
+		b.bits[ls&sm] &= ^(((1 << (ro - lo + 1)) - 1) << lo)
 	} else {
-		b.bits[ls&b.mask] &= ^uint64(0) >> (64 - lo)
+		b.bits[ls&sm] &= ^uint64(0) >> (64 - lo)
 		for i := ls + 1; i < rs; i++ {
-			b.bits[i&b.mask] = 0
+			b.bits[i&sm] = 0
 		}
-		b.bits[rs&b.mask] &= ^uint64(0) << (ro + 1)
+		b.bits[rs&sm] &= ^uint64(0) << (ro + 1)
 	}
 }
 
 func (b *Bitmap[T]) IsSet(val T) bool {
-	s, o := b.getSlotAndOffset(val)
-	return b.bits[s&b.mask]&(1<<o) != 0
+	sm, s, o := b.getSlotAndOffset(val)
+	return b.bits[s&sm]&(1<<o) != 0
 }
 
-func (b *Bitmap[T]) getSlotAndOffset(val T) (s, o int) {
-	s = int(val >> 6)   // slot
-	o = int(val & 0x3f) // offset
+func (b *Bitmap[T]) getSlotAndOffset(val T) (sm, s, o int) {
+	sm = len(b.bits) - 1 // slot mask
+	s = int(val >> 6)    // slot
+	o = int(val & 0x3f)  // offset
 	return
 }
 
-func (b *Bitmap[T]) getSlotsAndOffsets(min, max T) (ls, rs, lo, ro int) {
+func (b *Bitmap[T]) getSlotsAndOffsets(min, max T) (sm, ls, rs, lo, ro int) {
+	sm = len(b.bits) - 1 // slot mask
+
 	ls = int(min >> 6) // left slot
 	rs = int(max >> 6) // right slot
 

utils/bitmap_test.go

@@ -94,6 +94,31 @@ func TestBitmap(t *testing.T) {
 	require.Equal(t, e, b.bits)
 
 	// large range changes touch each word once
-	ls, rs, lo, ro := b.getSlotsAndOffsets(0, math.MaxUint32)
-	require.Equal(t, []int{0, 4, 0, 0}, []int{ls, rs, lo, ro})
+	sm, ls, rs, lo, ro := b.getSlotsAndOffsets(0, math.MaxUint32)
+	require.Equal(t, []int{3, 0, 4, 0, 0}, []int{sm, ls, rs, lo, ro})
+}
+
+func BenchmarkBitmap(b *testing.B) {
+	m := NewBitmap[uint64](128)
+	b.Run("SetRange", func(b *testing.B) {
+		for range b.N {
+			m.SetRange(100, 1100)
+		}
+	})
+	b.Run("Set", func(b *testing.B) {
+		for i := range b.N {
+			m.Set(uint64(i))
+		}
+	})
+	b.Run("IsSet/Set", func(b *testing.B) {
+		for i := range b.N {
+			_ = m.IsSet(uint64(i))
+			m.Set(uint64(i))
+		}
+	})
+	b.Run("GetAndSet", func(b *testing.B) {
+		for i := range b.N {
+			m.GetAndSet(uint64(i))
+		}
+	})
 }

utils/math.go

@@ -14,11 +14,15 @@
 
 package utils
 
-import "math"
+import (
+	"math"
+
+	"golang.org/x/exp/constraints"
+)
 
 func LogisticFunc(x0, L, k float64) func(x float64) float64 {
 	return func(x float64) float64 {
-		return L / (1 + math.Pow(math.E, -k*(x-x0)))
+		return L / (1 + math.Exp(-k*(x-x0)))
 	}
 }
 
@@ -27,3 +31,10 @@ func FastLogisticFunc(x0, L, k float64) func(x float64) float64 {
 		return L / 2 * (1 + k*(x-x0)/(1+math.Abs(k*(x-x0))))
 	}
 }
+
+func Abs[T constraints.Signed | constraints.Float](v T) T {
+	if v < 0 {
+		return -v
+	}
+	return v
+}

utils/math_test.go

@@ -0,0 +1,40 @@
+// Copyright 2023 LiveKit, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package utils
+
+import (
+	"testing"
+)
+
+func BenchmarkLogisticFunc(b *testing.B) {
+	b.Run("LogisticFunc", func(b *testing.B) {
+		l := LogisticFunc(0, 1, 1)
+		step := 4.0 / float64(b.N)
+		var i float64
+		for range b.N {
+			_ = l(i)
+			i += step
+		}
+	})
+	b.Run("FastLogisticFunc", func(b *testing.B) {
+		l := FastLogisticFunc(0, 1, 1)
+		step := 4.0 / float64(b.N)
+		var i float64
+		for range b.N {
+			_ = l(i)
+			i += step
+		}
+	})
+}

utils/timed_version.go

@@ -166,11 +166,7 @@ func (t TimedVersion) Value() (driver.Value, error) {
 		return nil, nil
 	}
 
-	ts, ticks := timedVersionComponents(t)
-	b := make([]byte, 0, 12)
-	b = binary.BigEndian.AppendUint64(b, uint64(ts))
-	b = binary.BigEndian.AppendUint32(b, uint32(ticks))
-	return b, nil
+	return t.MarshalBinary()
 }
 
 func (t *TimedVersion) Scan(src interface{}) (err error) {
@@ -180,9 +176,7 @@ func (t *TimedVersion) Scan(src interface{}) (err error) {
 		case 0:
 			*t = 0
 		case 12:
-			ts := int64(binary.BigEndian.Uint64(b))
-			ticks := int32(binary.BigEndian.Uint32(b[8:]))
-			*t = timedVersionFromComponents(ts, ticks)
+			t.UnmarshalBinary(b)
 		default:
 			return errors.New("(*TimedVersion).Scan: unsupported format")
 		}
@@ -194,6 +188,21 @@ func (t *TimedVersion) Scan(src interface{}) (err error) {
 	return nil
 }
 
+func (t TimedVersion) MarshalBinary() ([]byte, error) {
+	ts, ticks := timedVersionComponents(t)
+	b := make([]byte, 0, 12)
+	b = binary.BigEndian.AppendUint64(b, uint64(ts))
+	b = binary.BigEndian.AppendUint32(b, uint32(ticks))
+	return b, nil
+}
+
+func (t *TimedVersion) UnmarshalBinary(b []byte) error {
+	ts := int64(binary.BigEndian.Uint64(b))
+	ticks := int32(binary.BigEndian.Uint32(b[8:]))
+	*t = timedVersionFromComponents(ts, ticks)
+	return nil
+}
+
 func (t TimedVersion) GormDataType() string {
 	return "bytes"
 }