Kaynağa Gözat

Update buffer to use more modern go standards for string/byte conversion

Matt Bolt 1 yıl önce
ebeveyn
işleme
f9391932ac
2 değiştirilmiş dosya ile 61 ekleme ve 9 silme
  1. 4 9
      core/pkg/util/buffer.go
  2. 57 0
      core/pkg/util/buffer_test.go

+ 4 - 9
core/pkg/util/buffer.go

@@ -6,7 +6,6 @@ import (
 	"errors"
 	"io"
 	"math"
-	"reflect"
 	"unsafe"
 
 	"github.com/opencost/opencost/core/pkg/util/stringutil"
@@ -114,6 +113,7 @@ func (b *Buffer) WriteFloat64(i float64) {
 // WriteString writes the string's length as a uint16 followed by the string contents.
 func (b *Buffer) WriteString(i string) {
 	s := stringToBytes(i)
+
 	// string lengths are limited to uint16 - See ReadString()
 	if len(s) > math.MaxUint16 {
 		s = s[:math.MaxUint16]
@@ -401,7 +401,7 @@ func bytesToString(b []byte) string {
 	// cached string. If it does _not_ exist, then we use the passed func() string to allocate a new
 	// string and cache it. This will prevent us from allocating throw-away strings just to
 	// check our cache.
-	pinned := *(*string)(unsafe.Pointer(&b))
+	pinned := unsafe.String(unsafe.SliceData(b), len(b))
 
 	return stringutil.BankFunc(pinned, func() string {
 		return string(b)
@@ -409,11 +409,6 @@ func bytesToString(b []byte) string {
 }
 
 // Direct string to byte conversion that doesn't allocate.
-func stringToBytes(s string) (b []byte) {
-	strh := (*reflect.StringHeader)(unsafe.Pointer(&s))
-	sh := (*reflect.SliceHeader)(unsafe.Pointer(&b))
-	sh.Data = strh.Data
-	sh.Len = strh.Len
-	sh.Cap = strh.Len
-	return b
+func stringToBytes(s string) []byte {
+	return unsafe.Slice(unsafe.StringData(s), len(s))
 }

+ 57 - 0
core/pkg/util/buffer_test.go

@@ -4,7 +4,10 @@ import (
 	"bytes"
 	"math"
 	"math/rand"
+	"runtime"
+	"strings"
 	"testing"
+	"time"
 )
 
 func TestBufferReadWrite(t *testing.T) {
@@ -226,6 +229,60 @@ func generateRandomString(ln int) string {
 	return string(b)
 }
 
+func memMib() float64 {
+	var m runtime.MemStats
+	runtime.ReadMemStats(&m)
+	return float64(m.Alloc) / 1024.0 / 1024.0
+}
+
+func TestStringBytes(t *testing.T) {
+	baselineMem := memMib()
+
+	b := make([]byte, 10<<20)
+
+	afterMem := memMib()
+	delta := afterMem - baselineMem
+	t.Logf("Allocated %v MiB, Delta: %v MiB", afterMem, delta)
+
+	s := "Hello World!"
+	sl := b[512 : 512+len(s)]
+	copy(sl, stringToBytes(s))
+
+	afterMem = memMib()
+	delta = afterMem - baselineMem
+	t.Logf("Allocated %v MiB, Delta: %v MiB", afterMem, delta)
+
+	// this should pin the large backing array in memory, preventing it from being GC'd
+	newS := bytesToString(sl)
+
+	runtime.GC()
+	time.Sleep(time.Second)
+
+	afterMem = memMib()
+	delta = afterMem - baselineMem
+	t.Logf("S: %s, Allocated %v MiB, Delta: %v MiB", newS, afterMem, delta)
+
+	// copy the string into a new string and clear out pinned string
+	sCopy := strings.Clone(newS)
+	newS = ""
+
+	// Now that we've dropped the reference to the pinned backing array, it should be GC'd
+	runtime.GC()
+	time.Sleep(time.Second)
+
+	afterMem = memMib()
+	delta = afterMem - baselineMem
+	t.Logf("S: %s, Allocated %v MiB, Delta: %v MiB", sCopy, afterMem, delta)
+
+	if sCopy != s {
+		t.Errorf("Expected string to be %v, got %v", s, sCopy)
+	}
+
+	if delta > 0.5 {
+		t.Errorf("Expected memory delta to be less than 0.5 MiB, got %v MiB", delta)
+	}
+}
+
 func TestTooLargeStringTruncate(t *testing.T) {
 	normalStr := generateRandomString(100)
 	bigStr := generateRandomString(math.MaxUint16 + (math.MaxUint16 / 2))