feat: add Session.InternalState and Session.MergeEnv

node/session.go

@@ -44,7 +44,9 @@ type Session struct {
 	UID         string
 	Identifiers string
 	Connected   bool
-	Log         *log.Entry
+	// Could be used to store arbitrary data within a session
+	InternalState map[string]interface{}
+	Log           *log.Entry
 }
 
 // NewSession build a new Session struct from ws connetion and http request
@@ -95,6 +97,24 @@ func (s *Session) SetEnv(env *common.SessionEnv) {
 	s.env = env
 }
 
+// Merge connection and channel states into current env.
+// This method locks the state for writing (so, goroutine-safe)
+func (s *Session) MergeEnv(env *common.SessionEnv) {
+	s.smu.Lock()
+	defer s.smu.Unlock()
+
+	if env.ConnectionState != nil {
+		s.env.MergeConnectionState(env.ConnectionState)
+	}
+
+	if env.ChannelStates != nil {
+		states := *env.ChannelStates
+		for id, state := range states { // #nosec
+			s.env.MergeChannelState(id, &state)
+		}
+	}
+}
+
 // Serve enters a loop to read incoming data
 func (s *Session) Serve(callback func()) error {
 	go func() {

node/session_test.go

@@ -4,6 +4,7 @@ import (
 	"sync"
 	"testing"
 
+	"github.com/anycable/anycable-go/common"
 	"github.com/anycable/anycable-go/ws"
 	"github.com/stretchr/testify/assert"
 )
@@ -86,3 +87,59 @@ func TestSessionDisconnect(t *testing.T) {
 	_, err = session.conn.Read()
 	assert.Nil(t, err)
 }
+
+func TestMergeEnv(t *testing.T) {
+	node := NewMockNode()
+	session := NewMockSession("123", &node)
+
+	istate := map[string]map[string]string{
+		"test_channel": {
+			"foo": "bar",
+			"a":   "z",
+		},
+	}
+	cstate := map[string]string{"_s_": "id=42"}
+	origEnv := common.SessionEnv{ChannelStates: &istate, ConnectionState: &cstate}
+
+	session.SetEnv(&origEnv)
+
+	istate2 := map[string]map[string]string{
+		"test_channel": {
+			"foo": "baz",
+		},
+		"another_channel": {
+			"wasting": "time",
+		},
+	}
+
+	env := common.SessionEnv{ChannelStates: &istate2}
+
+	cstate2 := map[string]string{"red": "end of silence"}
+
+	env2 := common.SessionEnv{ConnectionState: &cstate2}
+
+	var wg sync.WaitGroup
+
+	wg.Add(2)
+
+	go func() {
+		session.MergeEnv(&env)
+		wg.Done()
+	}()
+
+	go func() {
+		session.MergeEnv(&env2)
+		wg.Done()
+	}()
+
+	wg.Wait()
+
+	assert.Equal(t, &origEnv, session.GetEnv())
+
+	assert.Equal(t, "id=42", origEnv.GetConnectionStateField("_s_"))
+	assert.Equal(t, "end of silence", origEnv.GetConnectionStateField("red"))
+
+	assert.Equal(t, "baz", origEnv.GetChannelStateField("test_channel", "foo"))
+	assert.Equal(t, "z", origEnv.GetChannelStateField("test_channel", "a"))
+	assert.Equal(t, "time", origEnv.GetChannelStateField("another_channel", "wasting"))
+}