[client] Fix state manager race conditions (#2890)

client/internal/statemanager/manager.go

@@ -74,15 +74,15 @@ func (m *Manager) Stop(ctx context.Context) error {
 	m.mu.Lock()
 	defer m.mu.Unlock()
 
-	if m.cancel != nil {
-		m.cancel()
+	if m.cancel == nil {
+		return nil
+	}
+	m.cancel()
 
-		select {
-		case <-ctx.Done():
-			return ctx.Err()
-		case <-m.done:
-			return nil
-		}
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	case <-m.done:
 	}
 
 	return nil
@@ -179,14 +179,18 @@ func (m *Manager) PersistState(ctx context.Context) error {
 		return nil
 	}
 
+	bs, err := marshalWithPanicRecovery(m.states)
+	if err != nil {
+		return fmt.Errorf("marshal states: %w", err)
+	}
+
 	ctx, cancel := context.WithTimeout(ctx, 5*time.Second)
 	defer cancel()
 
 	done := make(chan error, 1)
-
 	start := time.Now()
 	go func() {
-		done <- util.WriteJsonWithRestrictedPermission(ctx, m.filePath, m.states)
+		done <- util.WriteBytesWithRestrictedPermission(ctx, m.filePath, bs)
 	}()
 
 	select {
@@ -286,3 +290,19 @@ func (m *Manager) PerformCleanup() error {
 
 	return nberrors.FormatErrorOrNil(merr)
 }
+
+func marshalWithPanicRecovery(v any) ([]byte, error) {
+	var bs []byte
+	var err error
+
+	func() {
+		defer func() {
+			if r := recover(); r != nil {
+				err = fmt.Errorf("panic during marshal: %v", r)
+			}
+		}()
+		bs, err = json.Marshal(v)
+	}()
+
+	return bs, err
+}