From 7fa2dc4419bef0387da793a45c5b6c8037d3df78 Mon Sep 17 00:00:00 2001
From: tidwall <joshbaker77@gmail.com>
Date: Fri, 23 Sep 2022 12:42:39 -0700
Subject: [PATCH] Better FSET tests

Execute oom check immediately after setting maxmemory
---
 internal/server/config.go |  3 +++
 internal/server/server.go | 38 +++++++++++++++++---------------
 tests/keys_test.go        | 46 ++++++++++++++++++++++++---------------
 3 files changed, 51 insertions(+), 36 deletions(-)

diff --git a/internal/server/config.go b/internal/server/config.go
index b29cfc1d..7ec0bcad 100644
--- a/internal/server/config.go
+++ b/internal/server/config.go
@@ -427,6 +427,9 @@ func (s *Server) cmdConfigSet(msg *Message) (res resp.Value, err error) {
 	if err := s.config.setProperty(name, value, false); err != nil {
 		return NOMessage, err
 	}
+	if name == MaxMemory {
+		s.checkOutOfMemory()
+	}
 	return OKMessage(msg, start), nil
 }
 func (s *Server) cmdConfigRewrite(msg *Message) (res resp.Value, err error) {
diff --git a/internal/server/server.go b/internal/server/server.go
index 9af3f5f2..1730362e 100644
--- a/internal/server/server.go
+++ b/internal/server/server.go
@@ -623,28 +623,30 @@ func (s *Server) watchAutoGC() {
 	}
 }
 
+func (s *Server) checkOutOfMemory() {
+	if s.stopServer.on() {
+		return
+	}
+	oom := s.outOfMemory.on()
+	var mem runtime.MemStats
+	if s.config.maxMemory() == 0 {
+		if oom {
+			s.outOfMemory.set(false)
+		}
+		return
+	}
+	if oom {
+		runtime.GC()
+	}
+	runtime.ReadMemStats(&mem)
+	s.outOfMemory.set(int(mem.HeapAlloc) > s.config.maxMemory())
+}
+
 func (s *Server) watchOutOfMemory() {
 	t := time.NewTicker(time.Second * 2)
 	defer t.Stop()
-	var mem runtime.MemStats
 	for range t.C {
-		func() {
-			if s.stopServer.on() {
-				return
-			}
-			oom := s.outOfMemory.on()
-			if s.config.maxMemory() == 0 {
-				if oom {
-					s.outOfMemory.set(false)
-				}
-				return
-			}
-			if oom {
-				runtime.GC()
-			}
-			runtime.ReadMemStats(&mem)
-			s.outOfMemory.set(int(mem.HeapAlloc) > s.config.maxMemory())
-		}()
+		s.checkOutOfMemory()
 	}
 }
 
diff --git a/tests/keys_test.go b/tests/keys_test.go
index 3d4248cd..b2932816 100644
--- a/tests/keys_test.go
+++ b/tests/keys_test.go
@@ -160,24 +160,34 @@ func keys_EXPIRE_test(mc *mockServer) error {
 	)
 }
 func keys_FSET_test(mc *mockServer) error {
-	return mc.DoBatch([][]interface{}{
-		{"SET", "mykey", "myid", "HASH", "9my5xp7"}, {"OK"},
-		{"GET", "mykey", "myid", "WITHFIELDS", "HASH", 7}, {"[9my5xp7]"},
-		{"FSET", "mykey", "myid", "f1", 105.6}, {1},
-		{"GET", "mykey", "myid", "WITHFIELDS", "HASH", 7}, {"[9my5xp7 [f1 105.6]]"},
-		{"FSET", "mykey", "myid", "f1", 1.1, "f2", 2.2}, {2},
-		{"GET", "mykey", "myid", "WITHFIELDS", "HASH", 7}, {"[9my5xp7 [f1 1.1 f2 2.2]]"},
-		{"FSET", "mykey", "myid", "f1", 1.1, "f2", 22.22}, {1},
-		{"GET", "mykey", "myid", "WITHFIELDS", "HASH", 7}, {"[9my5xp7 [f1 1.1 f2 22.22]]"},
-		{"FSET", "mykey", "myid", "f1", 0}, {1},
-		{"GET", "mykey", "myid", "WITHFIELDS", "HASH", 7}, {"[9my5xp7 [f2 22.22]]"},
-		{"FSET", "mykey", "myid", "f2", 0}, {1},
-		{"GET", "mykey", "myid", "WITHFIELDS", "HASH", 7}, {"[9my5xp7]"},
-		{"FSET", "mykey", "myid2", "xx", "f1", 1.1, "f2", 2.2}, {0},
-		{"GET", "mykey", "myid2"}, {nil},
-		{"DEL", "mykey", "myid"}, {"1"},
-		{"GET", "mykey", "myid"}, {nil},
-	})
+	return mc.DoBatch(
+		Do("SET", "mykey", "myid", "HASH", "9my5xp7").OK(),
+		Do("GET", "mykey", "myid", "WITHFIELDS", "HASH", 7).Str("[9my5xp7]"),
+		Do("FSET", "mykey", "myid", "f1", 105.6).Str("1"),
+		Do("GET", "mykey", "myid", "WITHFIELDS", "HASH", 7).Str("[9my5xp7 [f1 105.6]]"),
+		Do("FSET", "mykey", "myid", "f1", 1.1, "f2", 2.2).Str("2"),
+		Do("GET", "mykey", "myid", "WITHFIELDS", "HASH", 7).Str("[9my5xp7 [f1 1.1 f2 2.2]]"),
+		Do("FSET", "mykey", "myid", "f1", 1.1, "f2", 22.22).Str("1"),
+		Do("GET", "mykey", "myid", "WITHFIELDS", "HASH", 7).Str("[9my5xp7 [f1 1.1 f2 22.22]]"),
+		Do("FSET", "mykey", "myid", "f1", 0).Str("1"),
+		Do("GET", "mykey", "myid", "WITHFIELDS", "HASH", 7).Str("[9my5xp7 [f2 22.22]]"),
+		Do("FSET", "mykey", "myid", "f2", 0).Str("1"),
+		Do("GET", "mykey", "myid", "WITHFIELDS", "HASH", 7).Str("[9my5xp7]"),
+		Do("FSET", "mykey", "myid2", "xx", "f1", 1.1, "f2", 2.2).Str("0"),
+		Do("GET", "mykey", "myid2").Str("<nil>"),
+		Do("DEL", "mykey", "myid").Str("1"),
+		Do("GET", "mykey", "myid").Str("<nil>"),
+		Do("SET", "mykey", "myid", "HASH", "9my5xp7").OK(),
+		Do("CONFIG", "SET", "maxmemory", "1").OK(),
+		Do("FSET", "mykey", "myid", "xx", "f1", 1.1, "f2", 2.2).Err(`OOM command not allowed when used memory > 'maxmemory'`),
+		Do("CONFIG", "SET", "maxmemory", "0").OK(),
+		Do("FSET", "mykey", "myid", "xx").Err("wrong number of arguments for 'fset' command"),
+		Do("FSET", "mykey", "myid", "f1", "a", "f2").Err("wrong number of arguments for 'fset' command"),
+		Do("FSET", "mykey", "myid", "z", "a").Err("invalid argument 'z'"),
+		Do("FSET", "mykey2", "myid", "a", "b").Err("key not found"),
+		Do("FSET", "mykey", "myid2", "a", "b").Err("id not found"),
+		Do("FSET", "mykey", "myid", "f2", 0).JSON().OK(),
+	)
 }
 func keys_GET_test(mc *mockServer) error {
 	return mc.DoBatch(