Add timeouts for lua scripts

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)
 	}