From 6631d8b00e8c5bafea15ea24415c5b55439193ee Mon Sep 17 00:00:00 2001
From: sangtran <tranquangsang12.7@gmail.com>
Date: Mon, 31 Aug 2026 23:51:37 +0700
Subject: bug: fix memory_state clear and guard window state empty with
 implement main.go for test

---
 state/memory_backend.go | 21 +++++++++++----------
 state/memory_state.go   | 18 +++++++++++++-----
 2 files changed, 24 insertions(+), 15 deletions(-)

(limited to 'state')

diff --git a/state/memory_backend.go b/state/memory_backend.go
index 6594657..f4043a1 100644
--- a/state/memory_backend.go
+++ b/state/memory_backend.go
@@ -2,6 +2,7 @@ package state
 
 import (
 	"context"
+
 	"flux/core"
 )
 
@@ -11,9 +12,9 @@ import (
 // This should be optimize with ZERO GC Scanning, Memory Alignment
 
 type MemoryStateBackend struct {
-	maxParallelism  int
-	currentKey      string
-	currentGroupKey int
+	maxParallelism   int
+	currentKey       string
+	currentGroupKey  int
 	currentNamespace any
 	// 4D structure StateName -> KeyGroup -> Key -> Namespace -> Value
 	states map[string]map[int]map[string]map[any]any
@@ -36,18 +37,18 @@ func (ms *MemoryStateBackend) SetCurrentKey(key string) error {
 	return nil
 }
 
-func (ms *MemoryStateBackend) GetValueState(ctx context.Context, key string) (core.ValueState[any], error) {
+func (ms *MemoryStateBackend) GetValueState(ctx context.Context, statename string) (core.ValueState[any], error) {
 	// Assign if state not existed
-
-	if _, ok := ms.states[key]; !ok {
-		ms.states[key] = make(map[int]map[string]map[any]any)
+	
+	if _, ok := ms.states[statename]; !ok {
+		// create state name map if not existed 
+		ms.states[statename] = make(map[int]map[string]map[any]any)
 	}
 
-	return &memoryStateValue{backend: ms, stateName: key}, nil
+	return &memoryStateValue{backend: ms, stateName: statename}, nil
 }
 
-func (ms *MemoryStateBackend) GetCurrentNamespace() any{
-
+func (ms *MemoryStateBackend) GetCurrentNamespace() any {
 	return ms.currentNamespace
 }
 
diff --git a/state/memory_state.go b/state/memory_state.go
index 95c7318..97ecad1 100644
--- a/state/memory_state.go
+++ b/state/memory_state.go
@@ -1,7 +1,7 @@
 package state
 
 import (
-"context"
+	"context"
 )
 
 type memoryStateValue struct {
@@ -13,7 +13,7 @@ func (m *memoryStateValue) Value(ctx context.Context) (any, error) {
 	kg := m.backend.currentGroupKey
 	key := m.backend.currentKey
 	ns := m.backend.currentNamespace
-	if kgMap, exist := m.backend.states[m.stateName][kg][key]; exist {
+	if kgMap, ok := m.backend.states[m.stateName][kg][key]; ok {
 		if v, hasKey := kgMap[ns]; hasKey {
 			return v, nil
 		}
@@ -26,10 +26,15 @@ func (m *memoryStateValue) Update(ctx context.Context, newVal any) (any, error)
 	kg := m.backend.currentGroupKey
 	key := m.backend.currentKey
 	ns := m.backend.currentNamespace
-	if _, exists := m.backend.states[m.stateName][kg][key]; !exists {
+
+	if _, ok := m.backend.states[m.stateName][kg]; !ok {
 		m.backend.states[m.stateName][kg] = make(map[string]map[any]any)
 	}
 
+	if _, ok := m.backend.states[m.stateName][kg][key]; !ok {
+		m.backend.states[m.stateName][kg][key] = make(map[any]any)
+	}
+
 	m.backend.states[m.stateName][kg][key][ns] = newVal
 	return newVal, nil
 }
@@ -37,8 +42,11 @@ func (m *memoryStateValue) Update(ctx context.Context, newVal any) (any, error)
 func (m *memoryStateValue) Clear(ctx context.Context) error {
 	kg := m.backend.currentGroupKey
 	key := m.backend.currentKey
-	if kgMap, exists := m.backend.states[m.stateName][kg]; exists {
-		delete(kgMap, key)
+	ns := m.backend.currentNamespace
+	if kgMap, ok := m.backend.states[m.stateName][kg]; ok {
+		if nsMap, ok := kgMap[key]; ok {
+			delete(nsMap, ns)
+		}
 	}
 	return nil
 }
-- 
cgit v1.3.1