Fixed up toProtocol compressed write flag

Signed-off-by: Jimmy Moore <[email protected]>
loopholelabs · Nov 25, 2024 · f84513d · f84513d
1 parent fd33b4f
commit f84513d
Show file tree

Hide file tree

Showing 4 changed files with 13 additions and 9 deletions.
diff --git a/pkg/storage/migrator/migrator_benchmark_test.go b/pkg/storage/migrator/migrator_benchmark_test.go
@@ -200,7 +200,7 @@ func BenchmarkMigrationPipe(mb *testing.B) {
 			destination := protocol.NewToProtocol(sourceDirtyRemote.Size(), 17, prSource)
 
 			if testconf.compress {
-				destination.CompressedWrites = true
+				destination.SetCompression(true)
 			}
 
 			err = destination.SendDevInfo("test", uint32(blockSize), "")

diff --git a/pkg/storage/protocol/protocol_transport_benchmark_test.go b/pkg/storage/protocol/protocol_transport_benchmark_test.go
@@ -106,7 +106,7 @@ func BenchmarkWriteAt(mb *testing.B) {
 func BenchmarkWriteAtComp(mb *testing.B) {
 	sourceToProtocol := setup(1)
 
-	sourceToProtocol.CompressedWrites = true
+	sourceToProtocol.SetCompression(true)
 
 	// Do some writes
 	buff := make([]byte, 256*1024)

diff --git a/pkg/storage/protocol/protocol_transport_test.go b/pkg/storage/protocol/protocol_transport_test.go
@@ -73,7 +73,7 @@ func TestProtocolWriteAtComp(t *testing.T) {
 
 	sourceToProtocol := NewToProtocol(uint64(size), 1, pr)
 
-	sourceToProtocol.CompressedWrites = true
+	sourceToProtocol.SetCompression(true)
 
 	storeFactory := func(di *packets.DevInfo) storage.Provider {
 		store = sources.NewMemoryStorage(int(di.Size))

diff --git a/pkg/storage/protocol/to_protocol.go b/pkg/storage/protocol/to_protocol.go
@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"crypto/sha256"
 	"encoding/binary"
+	"sync/atomic"
 
 	"github.com/loopholelabs/silo/pkg/storage"
 	"github.com/loopholelabs/silo/pkg/storage/protocol/packets"
@@ -14,16 +15,15 @@ type ToProtocol struct {
 	size             uint64
 	dev              uint32
 	protocol         Protocol
-	CompressedWrites bool
+	compressedWrites atomic.Bool
 	alternateSources []packets.AlternateSource
 }
 
 func NewToProtocol(size uint64, deviceID uint32, p Protocol) *ToProtocol {
 	return &ToProtocol{
-		size:             size,
-		dev:              deviceID,
-		protocol:         p,
-		CompressedWrites: false,
+		size:     size,
+		dev:      deviceID,
+		protocol: p,
 	}
 }
 
@@ -39,6 +39,10 @@ func (i *ToProtocol) SendSiloEvent(eventType storage.EventType, eventData storag
 	return nil
 }
 
+func (i *ToProtocol) SetCompression(compressed bool) {
+	i.compressedWrites.Store(compressed)
+}
+
 func (i *ToProtocol) SendEvent(e *packets.Event) error {
 	b := packets.EncodeEvent(e)
 	id, err := i.protocol.SendPacket(i.dev, IDPickAny, b, UrgencyUrgent)
@@ -150,7 +154,7 @@ func (i *ToProtocol) WriteAt(buffer []byte, offset int64) (int, error) {
 	}
 
 	if !dontSendData {
-		if i.CompressedWrites {
+		if i.compressedWrites.Load() {
 			data := packets.EncodeWriteAtComp(offset, buffer)
 			id, err = i.protocol.SendPacket(i.dev, IDPickAny, data, UrgencyNormal)
 		} else {