Merge pull request #444 from rshura/script-timeouts

Script timeouts

internal/deadline/deadline.go

@@ -13,6 +13,7 @@ func New(deadline time.Time) *Deadline {
 	return &Deadline{unixNano: deadline.UnixNano()}
 }
 
+// Empty deadline does nothing, just a place holder for future updates
 func Empty() *Deadline {
 	return &Deadline{}
 }
@@ -38,3 +39,8 @@ func (deadline *Deadline) Check() {
 func (deadline *Deadline) Hit() bool {
 	return deadline.hit
 }
+
+// GetDeadlineTime returns the time object for the deadline, and an "empty" boolean
+func (deadline *Deadline) GetDeadlineTime() (time.Time, bool) {
+	return time.Unix(0, deadline.unixNano), deadline.unixNano == 0
+}

internal/server/scripts.go

@@ -2,6 +2,7 @@ package server
 
 import (
 	"bytes"
+	"context"
 	"crypto/sha1"
 	"encoding/hex"
 	"encoding/json"
@@ -391,6 +392,13 @@ func (c *Server) cmdEvalUnified(scriptIsSha bool, msg *Message) (res resp.Value,
 	if err != nil {
 		return
 	}
+	deadline, empty := msg.Deadline.GetDeadlineTime()
+	if !empty {
+		ctx, cancel := context.WithDeadline(context.Background(), deadline)
+		defer cancel()
+		luaState.SetContext(ctx)
+		defer luaState.RemoveContext()
+	}
 	defer c.luapool.Put(luaState)
 
 	keysTbl := luaState.CreateTable(int(numkeys), 0)
@@ -454,6 +462,9 @@ func (c *Server) cmdEvalUnified(scriptIsSha bool, msg *Message) (res resp.Value,
 			"EVAL_CMD": lua.LNil,
 		})
 	if err := luaState.PCall(0, 1, nil); err != nil {
+		if strings.Contains(err.Error(), "context deadline exceeded") {
+			msg.Deadline.Check()
+		}
 		log.Debugf("%v", err.Error())
 		return NOMessage, makeSafeErr(err)
 	}

tests/timeout_test.go

@@ -13,6 +13,7 @@ func subTestTimeout(t *testing.T, mc *mockServer) {
 	runStep(t, mc, "session set/unset", timeout_session_set_unset_test)
 	runStep(t, mc, "session spatial", timeout_session_spatial_test)
 	runStep(t, mc, "session search", timeout_session_search_test)
+	runStep(t, mc, "session scripts", timeout_session_scripts_test)
 	runStep(t, mc, "command spatial", timeout_command_spatial_test)
 	runStep(t, mc, "command search", timeout_command_search_test)
 }
@@ -112,3 +113,35 @@ func timeout_command_search_test(mc *mockServer) (err error) {
 		{"SEARCH", "mykey", "TIMEOUT", "0.000001", "MATCH", "val:*", "COUNT"}, {"ERR timeout"},
 	})
 }
+
+func timeout_session_scripts_test(mc *mockServer) (err error) {
+	script := `
+		local clock = os.clock
+		local function sleep(n)
+			local t0 = clock()
+			while clock() - t0 <= n do end
+		end
+		sleep(0.5)
+	`
+	sha := "e3ce9449853a622327f30c727a6e086ccd91d9d4"
+
+	return mc.DoBatch([][]interface{}{
+		{"SCRIPT LOAD", script}, {sha},
+
+		{"EVALSHA", sha, 0}, {nil},
+		{"EVALROSHA", sha, 0}, {nil},
+		{"EVALNASHA", sha, 0}, {nil},
+
+		{"TIMEOUT", "0.1"}, {"OK"},
+
+		{"EVALSHA", sha, 0}, {"ERR timeout"},
+		{"EVALROSHA", sha, 0}, {"ERR timeout"},
+		{"EVALNASHA", sha, 0}, {"ERR timeout"},
+
+		{"TIMEOUT", "0.9"}, {"OK"},
+
+		{"EVALSHA", sha, 0}, {nil},
+		{"EVALROSHA", sha, 0}, {nil},
+		{"EVALNASHA", sha, 0}, {nil},
+	})
+}