Fix concurrent map races in Statelet, collector, and session cache

service/executor/expand/data_unit.go

@@ -173,6 +173,8 @@ func (c *DataUnit) Next() (interface{}, error) {
 }
 
 func (c *DataUnit) ensureSliceIndex() {
+	c.mu.Lock()
+	defer c.mu.Unlock()
 	if c.sliceIndex != nil {
 		return
 	}
@@ -181,6 +183,11 @@ func (c *DataUnit) ensureSliceIndex() {
 }
 
 func (c *DataUnit) xunsafeSlice(valueType reflect.Type) *xunsafe.Slice {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	if c.sliceIndex == nil {
+		c.sliceIndex = map[reflect.Type]*xunsafe.Slice{}
+	}
 	slice, ok := c.sliceIndex[valueType]
 	if !ok {
 		slice = xunsafe.NewSlice(reflect.SliceOf(valueType))

service/executor/expand/data_unit_test.go

@@ -0,0 +1,35 @@
+package expand
+
+import (
+	"reflect"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+)
+
+// TestDataUnitConcurrentSliceIndex reproduces concurrent map write on
+// DataUnit.sliceIndex. xunsafeSlice writes the map with no lock.
+func TestDataUnitConcurrentSliceIndex(t *testing.T) {
+	deadline := time.Now().Add(3 * time.Second)
+	unit := &DataUnit{}
+	unit.ensureSliceIndex()
+	var next atomic.Uint64
+	var waitGroup sync.WaitGroup
+	waitGroup.Add(5)
+
+	run := func(fn func()) {
+		defer waitGroup.Done()
+		for time.Now().Before(deadline) {
+			fn()
+		}
+	}
+
+	for worker := 0; worker < 5; worker++ {
+		go run(func() {
+			n := int(next.Add(1)%64) + 1
+			_ = unit.xunsafeSlice(reflect.ArrayOf(n, reflect.TypeOf(byte(0))))
+		})
+	}
+	waitGroup.Wait()
+}

service/session/cache.go

@@ -24,10 +24,18 @@ func (c *cache) lookup(parameter *state.Parameter) (interface{}, bool) {
 }
 
 func (s *Session) MarshalJSON() ([]byte, error) {
-	return json.Marshal(s.cache.values)
+	s.cache.RWMutex.RLock()
+	snapshot := make(map[string]interface{}, len(s.cache.values))
+	for key, value := range s.cache.values {
+		snapshot[key] = value
+	}
+	s.cache.RWMutex.RUnlock()
+	return json.Marshal(snapshot)
 }
 
 func (s *Session) Unmarshal(parameters state.Parameters, data []byte) error {
+	s.cache.RWMutex.Lock()
+	defer s.cache.RWMutex.Unlock()
 	err := json.Unmarshal(data, &s.cache.values)
 	if err != nil {
 		return err

service/session/cache_test.go

@@ -0,0 +1,39 @@
+package session
+
+import (
+	"fmt"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/viant/datly/view/state"
+)
+
+// TestSessionMarshalJSONConcurrentCachePut reproduces concurrent map
+// iteration and write on session cache.values. MarshalJSON ranges the map
+// while put writes it.
+func TestSessionMarshalJSONConcurrentCachePut(t *testing.T) {
+	deadline := time.Now().Add(3 * time.Second)
+	s := &Session{cache: newCache()}
+	var next atomic.Uint64
+	var waitGroup sync.WaitGroup
+	waitGroup.Add(5)
+
+	run := func(fn func()) {
+		defer waitGroup.Done()
+		for time.Now().Before(deadline) {
+			fn()
+		}
+	}
+
+	go run(func() {
+		s.cache.put(&state.Parameter{Name: fmt.Sprintf("p%d", next.Add(1))}, 1)
+	})
+	for worker := 0; worker < 4; worker++ {
+		go run(func() {
+			_, _ = s.MarshalJSON()
+		})
+	}
+	waitGroup.Wait()
+}

view/cache.go

@@ -1002,6 +1002,8 @@ func (c *Cache) applyWarmupFieldNames(selector *Statelet, fieldNames []string) {
 	if selector == nil || c.owner == nil || len(fieldNames) == 0 {
 		return
 	}
+	selector.columnNamesMu.Lock()
+	defer selector.columnNamesMu.Unlock()
 	if selector._columnNames == nil {
 		selector._columnNames = map[string]bool{}
 	}
@@ -1019,7 +1021,7 @@ func (c *Cache) applyWarmupFieldNames(selector *Statelet, fieldNames []string) {
 		if outputName == "" {
 			outputName = columnName
 		}
-		if selector.Has(columnName) || selector.Has(outputName) {
+		if selector._columnNames[columnName] || selector._columnNames[outputName] {
 			continue
 		}
 		selector._columnNames[columnName] = true