Implement a watchdog timer for scripts that take too long to execute (bug 5837, r=fyren).

--HG--
extra : rebase_source : ffacb38457eca581660ce8f15c444ad828b7fedd

sourcepawn/jit/AMBuilder

@@ -12,6 +12,10 @@ compiler['CXXINCLUDES'].append(os.path.join(base, 'public', 'jit'))
 compiler['CXXINCLUDES'].append(os.path.join(base, 'public', 'jit', 'x86'))
 compiler['CXXINCLUDES'].append(os.path.join(base, 'knight', 'shared'))
 
+if AMBuild.target['platform'] == 'linux':
+	compiler['POSTLINKFLAGS'].append('-lpthread')
+	compiler['POSTLINKFLAGS'].append('-lrt')
+
 extension = AMBuild.AddJob('sourcepawn.jit.x86')
 binary = Cpp.LibraryBuilder('sourcepawn.jit.x86', AMBuild, extension, compiler)
 binary.AddSourceFiles('sourcepawn/jit', [
@@ -23,6 +27,7 @@ binary.AddSourceFiles('sourcepawn/jit', [
 	'sp_vm_engine.cpp',
 	'sp_vm_function.cpp',
 	'opcodes.cpp',
+	'watchdog_timer.cpp',
 	'x86/jit_x86.cpp',
 	'zlib/adler32.c',
 	'zlib/compress.c',

sourcepawn/jit/BaseRuntime.cpp

@@ -37,10 +37,21 @@ BaseRuntime::BaseRuntime()
   
   memset(m_CodeHash, 0, sizeof(m_CodeHash));
   memset(m_DataHash, 0, sizeof(m_DataHash));
+
+  ke::AutoLock lock(g_Jit.Mutex());
+  g_Jit.RegisterRuntime(this);
 }
 
 BaseRuntime::~BaseRuntime()
 {
+  // The watchdog thread takes the global JIT lock while it patches all
+  // runtimes. It is not enough to ensure that the unlinking of the runtime is
+  // protected; we cannot delete functions or code while the watchdog might be
+  // executing. Therefore, the entire destructor is guarded.
+  ke::AutoLock lock(g_Jit.Mutex());
+
+  g_Jit.DeregisterRuntime(this);
+
   for (uint32_t i = 0; i < m_plugin.num_publics; i++)
     delete m_PubFuncs[i];
   delete [] m_PubFuncs;

sourcepawn/jit/BaseRuntime.h

@@ -4,6 +4,7 @@
 
 #include <sp_vm_api.h>
 #include <ke_vector.h>
+#include <ke_inline_list.h>
 #include "jit_shared.h"
 #include "sp_vm_function.h"
 
@@ -33,7 +34,9 @@ struct floattbl_t
 };
 
 /* Jit wants fast access to this so we expose things as public */
-class BaseRuntime : public SourcePawn::IPluginRuntime
+class BaseRuntime
+  : public SourcePawn::IPluginRuntime,
+    public InlineListNode<BaseRuntime>
 {
  public:
   BaseRuntime();
@@ -73,6 +76,13 @@ class BaseRuntime : public SourcePawn::IPluginRuntime
     return &m_plugin;
   }
 
+  size_t NumJitFunctions() const {
+    return m_JitFunctions.length();
+  }
+  JitFunction *GetJitFunction(size_t i) const {
+    return m_JitFunctions[i];
+  }
+
  private:
   void SetupFloatNativeRemapping();
 
@@ -103,3 +113,4 @@ class BaseRuntime : public SourcePawn::IPluginRuntime
 };
 
 #endif //_INCLUDE_SOURCEPAWN_JIT_RUNTIME_H_
+

sourcepawn/jit/Makefile.shell

@@ -19,6 +19,7 @@ OBJECTS = dll_exports.cpp 	\
 	BaseRuntime.cpp		\
 	jit_function.cpp	\
 	opcodes.cpp		\
+	watchdog_timer.cpp \
 	md5/md5.cpp			\
 	zlib/adler32.c		\
 	zlib/compress.c		\
@@ -47,7 +48,7 @@ CXX_GCC4_FLAGS = -fvisibility-inlines-hidden
 CXX = c++
 CC = cc
 
-LINK = -m32 -lm
+LINK = -m32 -lm -lpthread -lrt
 
 INCLUDE = -I. -I.. -I$(SMSDK)/public -I$(SMSDK)/public/jit -I$(SMSDK)/public/jit/x86 \
 	-I$(SMSDK)/public/sourcepawn -I$(MMSOURCE17)/core/sourcehook -I$(SMSDK)/knight/shared -Ix86

sourcepawn/jit/dll_exports.cpp

@@ -100,7 +100,7 @@ public:
 		while (error->GetTraceInfo(&stk_info))
 		{
 			fprintf(stderr,
-			    "   [%d]  Line %d, %s::%s()",
+			    "   [%d]  Line %d, %s::%s()\n",
 				i++,
 				stk_info.line,
 				stk_info.filename,
@@ -221,6 +221,7 @@ int main(int argc, char **argv)
 
 	ShellDebugListener debug;
 	g_engine1.SetDebugListener(&debug);
+	g_engine2.InstallWatchdogTimer(5000);
 
 	int errcode = Execute(argv[1]);
 

sourcepawn/jit/engine2.cpp

@@ -1,3 +1,4 @@
+// vim: set ts=4 sw=4 tw=99 noet:
 #include <sourcemod_version.h>
 #include <stdio.h>
 #include <stdlib.h>
@@ -7,6 +8,7 @@
 #include "zlib/zlib.h"
 #include "BaseRuntime.h"
 #include "sp_vm_engine.h"
+#include "watchdog_timer.h"
 
 using namespace SourcePawn;
 
@@ -210,3 +212,9 @@ IPluginRuntime *SourcePawnEngine2::CreateEmptyRuntime(const char *name, uint32_t
 	
 	return rt;
 }
+
+bool SourcePawnEngine2::InstallWatchdogTimer(size_t timeout_ms)
+{
+	return g_WatchdogTimer.Initialize(timeout_ms);
+}
+

sourcepawn/jit/engine2.h

@@ -26,6 +26,7 @@ namespace SourcePawn
 		bool Initialize();
 		void Shutdown();
 		IPluginRuntime *CreateEmptyRuntime(const char *name, uint32_t memory);
+		bool InstallWatchdogTimer(size_t timeout_ms);
 	public:
 		IProfiler *GetProfiler();
 	private:

sourcepawn/jit/jit_function.cpp

@@ -1,23 +1,31 @@
+// vim: set ts=8 ts=2 sw=2 tw=99 et:
 #include "jit_function.h"
 #include "sp_vm_engine.h"
 #include "jit_x86.h"
 
-JitFunction::JitFunction(void *entry_addr, cell_t pcode_offs)
-: m_pEntryAddr(entry_addr), m_PcodeOffs(pcode_offs)
+JitFunction::JitFunction(void *entry_addr, cell_t pcode_offs, LoopEdge *edges, uint32_t nedges)
+  : m_pEntryAddr(entry_addr),
+    m_PcodeOffs(pcode_offs),
+    edges_(edges),
+    nedges_(nedges)
 {
 }
 
 JitFunction::~JitFunction()
 {
-	g_Jit.FreeCode(m_pEntryAddr);
+  delete [] edges_;
+  g_Jit.FreeCode(m_pEntryAddr);
 }
 
-void *JitFunction::GetEntryAddress()
+void *
+JitFunction::GetEntryAddress() const
 {
-	return m_pEntryAddr;
+  return m_pEntryAddr;
 }
 
-cell_t JitFunction::GetPCodeAddress()
+cell_t
+JitFunction::GetPCodeAddress() const
 {
-	return m_PcodeOffs;
+  return m_PcodeOffs;
 }
+

sourcepawn/jit/jit_function.h

@@ -1,19 +1,38 @@
+// vim: set ts=8 sts=2 sw=2 tw=99 et:
 #ifndef _INCLUDE_SOURCEPAWN_JIT2_FUNCTION_H_
 #define _INCLUDE_SOURCEPAWN_JIT2_FUNCTION_H_
 
 #include <sp_vm_types.h>
+#include <ke_vector.h>
+
+struct LoopEdge
+{
+  uint32_t offset;
+  int32_t disp32;
+};
 
 class JitFunction
 {
-public:
-	JitFunction(void *entry_addr, cell_t pcode_offs);
-	~JitFunction();
-public:
-	void *GetEntryAddress();
-	cell_t GetPCodeAddress();
-private:
-	void *m_pEntryAddr;
-	cell_t m_PcodeOffs;
+ public:
+  JitFunction(void *entry_addr, cell_t pcode_offs, LoopEdge *edges, uint32_t nedges);
+  ~JitFunction();
+
+ public:
+  void *GetEntryAddress() const;
+  cell_t GetPCodeAddress() const;
+  uint32_t NumLoopEdges() const {
+    return nedges_;
+  }
+  const LoopEdge &GetLoopEdge(size_t i) const {
+    assert(i < nedges_);
+    return edges_[i];
+  }
+
+ private:
+  void *m_pEntryAddr;
+  cell_t m_PcodeOffs;
+  LoopEdge *edges_;
+  uint32_t nedges_;
 };
 
 #endif //_INCLUDE_SOURCEPAWN_JIT2_FUNCTION_H_

sourcepawn/jit/sp_vm_basecontext.cpp

@@ -36,6 +36,7 @@
 #include "sp_vm_api.h"
 #include "sp_vm_basecontext.h"
 #include "sp_vm_engine.h"
+#include "watchdog_timer.h"
 #include "x86/jit_x86.h"
 
 using namespace SourcePawn;
@@ -554,6 +555,9 @@ int BaseContext::Execute2(IPluginFunction *function, const cell_t *params, unsig
 
 	fnid = function->GetFunctionID();
 
+	if (!g_WatchdogTimer.HandleInterrupt())
+		return SP_ERROR_TIMEOUT;
+
 	if (fnid & 1)
 	{	
 		public_id = fnid >> 1;
@@ -674,6 +678,9 @@ int BaseContext::Execute2(IPluginFunction *function, const cell_t *params, unsig
 		}
 	}
 
+	if (ir == SP_ERROR_TIMEOUT)
+		g_WatchdogTimer.NotifyTimeoutReceived();
+
 	if (ir != SP_ERROR_NONE)
 	{
 		g_engine1.ReportError(m_pRuntime, ir, m_MsgCache, save_rp);

sourcepawn/jit/sp_vm_engine.cpp

@@ -1,5 +1,5 @@
 /**
- * vim: set ts=4 :
+ * vim: set ts=4 sw=4 tw=99 et:
  * =============================================================================
  * SourcePawn
  * Copyright (C) 2004-2008 AlliedModders LLC.  All rights reserved.
@@ -58,7 +58,7 @@ SourcePawnEngine g_engine1;
 
 using namespace SourcePawn;
 
-#define ERROR_MESSAGE_MAX		29
+#define ERROR_MESSAGE_MAX		30
 static const char *g_ErrorMsgTable[] = 
 {
 	NULL,
@@ -90,7 +90,8 @@ static const char *g_ErrorMsgTable[] =
 	"Plugin format is too old",
 	"Plugin format is too new",
     "Out of memory",
-    "Integer overflow"
+    "Integer overflow",
+    "Script execution timed out"
 };
 
 const char *SourcePawnEngine::GetErrorString(int error)

sourcepawn/jit/watchdog_timer.cpp

@@ -0,0 +1,164 @@
+// vim: set ts=8 sts=2 sw=2 tw=99 et:
+//
+// This file is part of SourcePawn.
+// 
+// SourcePawn is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+// 
+// SourcePawn is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+// 
+// You should have received a copy of the GNU General Public License
+// along with SourcePawn.  If not, see <http://www.gnu.org/licenses/>.
+#include "watchdog_timer.h"
+#include "x86/jit_x86.h"
+#include <string.h>
+
+WatchdogTimer g_WatchdogTimer;
+
+WatchdogTimer::WatchdogTimer()
+  : terminate_(false),
+    mainthread_(ke::GetCurrentThreadId()),
+    last_frame_id_(0),
+    second_timeout_(false),
+    timedout_(false)
+{
+}
+
+WatchdogTimer::~WatchdogTimer()
+{
+  Shutdown();
+}
+
+bool
+WatchdogTimer::Initialize(size_t timeout_ms)
+{
+  if (thread_)
+    return false;
+
+  timeout_ms_ = timeout_ms;
+
+  thread_ = new ke::Thread(this, "SM Watchdog");
+  if (!thread_->Succeeded())
+    return false;
+
+  return true;
+}
+
+void
+WatchdogTimer::Shutdown()
+{
+  if (terminate_ || !thread_)
+    return;
+
+  {
+    ke::AutoLock lock(&cv_);
+    terminate_ = true;
+    cv_.Notify();
+  }
+  thread_->Join();
+  thread_ = NULL;
+}
+
+void
+WatchdogTimer::Run()
+{
+  ke::AutoLock lock(&cv_);
+
+  // Initialize the frame id, so we don't have to wait longer on startup.
+  last_frame_id_ = g_Jit.FrameId();
+
+  while (!terminate_) {
+    ke::WaitResult rv = cv_.Wait(timeout_ms_ / 2);
+    if (terminate_)
+      return;
+
+    if (rv == ke::Wait_Error)
+      return;
+
+    if (rv != ke::Wait_Timeout)
+      continue;
+
+    // We reached a timeout. If the current frame is not equal to the last
+    // frame, then we assume the server is still moving enough to process
+    // frames. We also make sure JIT code is actually on the stack, since
+    // our concept of frames won't move if JIT code is not running.
+    //
+    // Note that it's okay if these two race: it's just a heuristic, and
+    // worst case, we'll miss something that might have timed out but
+    // ended up resuming.
+    uintptr_t frame_id = g_Jit.FrameId();
+    if (frame_id != last_frame_id_ || !g_Jit.RunningCode()) {
+      last_frame_id_ = frame_id;
+      second_timeout_ = false;
+      continue;
+    }
+
+    // We woke up with the same frame. We might be timing out. Wait another
+    // time to see if so.
+    if (!second_timeout_) {
+      second_timeout_ = true;
+      continue;
+    }
+
+    {
+      // Prevent the JIT from linking or destroying runtimes and functions.
+      ke::AutoLock lock(g_Jit.Mutex());
+
+      // Set the timeout notification bit. If this is detected before any patched
+      // JIT backedges are reached, the main thread will attempt to acquire the
+      // monitor lock, and block until we call Wait().
+      timedout_ = true;
+    
+      // Patch all jumps. This can race with the main thread's execution since
+      // all code writes are 32-bit integer instruction operands, which are
+      // guaranteed to be atomic on x86.
+      g_Jit.PatchAllJumpsForTimeout();
+    }
+
+    // The JIT will be free to compile new functions while we wait, but it will
+    // see the timeout bit set above and immediately bail out.
+    cv_.Wait();
+
+    second_timeout_ = false;
+
+    // Reset the last frame ID to something unlikely to be chosen again soon.
+    last_frame_id_--;
+
+    // It's possible that Shutdown() raced with the timeout, and if so, we
+    // must leave now to prevent a deadlock.
+    if (terminate_)
+      return;
+  }
+}
+
+bool
+WatchdogTimer::NotifyTimeoutReceived()
+{
+  // We are guaranteed that the watchdog thread is waiting for our
+  // notification, and is therefore blocked. We take the JIT lock
+  // anyway for sanity.
+  {
+    ke::AutoLock lock(g_Jit.Mutex());
+    g_Jit.UnpatchAllJumpsFromTimeout();
+  }
+
+  timedout_ = false;
+
+  // Wake up the watchdog thread, it's okay to keep processing now.
+  ke::AutoLock lock(&cv_);
+  cv_.Notify();
+  return false;
+}
+
+bool
+WatchdogTimer::HandleInterrupt()
+{
+  if (timedout_)
+    return NotifyTimeoutReceived();
+  return true;
+}

sourcepawn/jit/watchdog_timer.h

@@ -0,0 +1,61 @@
+// vim: set ts=8 sts=2 sw=2 tw=99 et:
+//
+// This file is part of SourcePawn.
+// 
+// SourcePawn is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+// 
+// SourcePawn is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+// 
+// You should have received a copy of the GNU General Public License
+// along with SourcePawn.  If not, see <http://www.gnu.org/licenses/>.
+#ifndef _include_sourcepawn_watchdog_timer_posix_h_
+#define _include_sourcepawn_watchdog_timer_posix_h_
+
+#include <stddef.h>
+#include <stdint.h>
+#include <ke_thread_utils.h>
+
+typedef bool (*WatchdogCallback)();
+
+class WatchdogTimer : public ke::IRunnable
+{
+ public:
+  WatchdogTimer();
+  ~WatchdogTimer();
+
+  bool Initialize(size_t timeout_ms);
+  void Shutdown();
+
+  // Called from main thread.
+  bool NotifyTimeoutReceived();
+  bool HandleInterrupt();
+
+ private:
+  // Watchdog thread.
+  void Run();
+
+ private:
+  bool terminate_;
+  size_t timeout_ms_;
+  ke::ThreadId mainthread_;
+
+  ke::AutoPtr<ke::Thread> thread_;
+  ke::ConditionVariable cv_;
+
+  // Accessed only on the watchdog thread.
+  uintptr_t last_frame_id_;
+  bool second_timeout_;
+
+  // Accessed only on the main thread.
+  bool timedout_;
+};
+
+extern WatchdogTimer g_WatchdogTimer;
+
+#endif // _include_sourcepawn_watchdog_timer_posix_h_

sourcepawn/jit/x86/jit_x86.cpp

@@ -37,6 +37,7 @@
 #include "../engine2.h"
 #include "../BaseRuntime.h"
 #include "../sp_vm_basecontext.h"
+#include "watchdog_timer.h"
 
 using namespace Knight;
 
@@ -50,7 +51,7 @@ JITX86 g_Jit;
 KeCodeCache *g_pCodeCache = NULL;
 ISourcePawnEngine *engine = &g_engine1;
 
-static inline void *
+static inline uint8_t *
 LinkCode(AssemblerX86 &masm)
 {
   if (masm.outOfMemory())
@@ -61,7 +62,7 @@ LinkCode(AssemblerX86 &masm)
     return NULL;
 
   masm.emitToExecutableMemory(code);
-  return code;
+  return reinterpret_cast<uint8_t *>(code);
 }
 
 static inline ConditionCode
@@ -249,7 +250,7 @@ GenerateArray(BaseRuntime *rt, uint32_t argc, cell_t *argv, int autozero)
     return err;
 
   cell_t *base = reinterpret_cast<cell_t *>(rt->plugin()->memory + ctx->hp);
-  cell_t offs = GenerateArrayIndirectionVectors(base, argv, argc, autozero);
+  cell_t offs = GenerateArrayIndirectionVectors(base, argv, argc, !!autozero);
   assert(size_t(offs) == cells);
 
   argv[argc - 1] = ctx->hp;
@@ -314,6 +315,12 @@ GetFunctionName(const sp_plugin_t *plugin, uint32_t offs)
 static int
 CompileFromThunk(BaseRuntime *runtime, cell_t pcode_offs, void **addrp, char *pc)
 {
+  // If the watchdog timer has declared a timeout, we must process it now,
+  // and possibly refuse to compile, since otherwise we will compile a
+  // function that is not patched for timeouts.
+  if (!g_WatchdogTimer.HandleInterrupt())
+    return SP_ERROR_TIMEOUT;
+
   JitFunction *fn = runtime->GetJittedFunctionByOffset(pcode_offs);
   if (!fn) {
     int err;
@@ -452,13 +459,19 @@ Compiler::emit(int *errp)
   emitCallThunks();
   emitErrorPaths();
 
-  void *code = LinkCode(masm);
+  uint8_t *code = LinkCode(masm);
   if (!code) {
     *errp = SP_ERROR_OUT_OF_MEMORY;
     return NULL;
   }
 
-  return new JitFunction(code, pcode_start_);
+  LoopEdge *edges = new LoopEdge[backward_jumps_.length()];
+  for (size_t i = 0; i < backward_jumps_.length(); i++) {
+    edges[i].offset = backward_jumps_[i];
+    edges[i].disp32 = *reinterpret_cast<int32_t *>(code + edges[i].offset - 4);
+  }
+
+  return new JitFunction(code, pcode_start_, edges, backward_jumps_.length());
 }
 
 bool
@@ -1228,7 +1241,12 @@ Compiler::emitOp(OPCODE op)
       Label *target = labelAt(readCell());
       if (!target)
         return false;
-      __ jmp(target);
+      if (target->bound()) {
+        __ jmp32(target);
+        backward_jumps_.append(masm.pc());
+      } else {
+        __ jmp(target);
+      }
       break;
     }
 
@@ -1240,7 +1258,12 @@ Compiler::emitOp(OPCODE op)
       if (!target)
         return false;
       __ testl(pri, pri);
-      __ j(cc, target);
+      if (target->bound()) {
+        __ j32(cc, target);
+        backward_jumps_.append(masm.pc());
+      } else {
+        __ j(cc, target);
+      }
       break;
     }
 
@@ -1256,7 +1279,12 @@ Compiler::emitOp(OPCODE op)
         return false;
       ConditionCode cc = OpToCondition(op);
       __ cmpl(pri, alt);
-      __ j(cc, target);
+      if (target->bound()) {
+        __ j32(cc, target);
+        backward_jumps_.append(masm.pc());
+      } else {
+        __ j(cc, target);
+      }
       break;
     }
 
@@ -1755,7 +1783,7 @@ Compiler::emitErrorPaths()
 typedef int (*JIT_EXECUTE)(sp_context_t *ctx, uint8_t *memory, void *code);
 
 static void *
-GenerateEntry(void **retp)
+GenerateEntry(void **retp, void **timeoutp)
 {
   AssemblerX86 masm;
 
@@ -1813,11 +1841,17 @@ GenerateEntry(void **retp)
   __ movl(ecx, Operand(ebp, 8 + 4 * 0)); // ret-path expects ecx = ctx
   __ jmp(&ret);
 
+  Label timeout;
+  __ bind(&timeout);
+  __ movl(eax, SP_ERROR_TIMEOUT);
+  __ jmp(&error);
+
   void *code = LinkCode(masm);
   if (!code)
     return NULL;
 
   *retp = reinterpret_cast<uint8_t *>(code) + error.offset();
+  *timeoutp = reinterpret_cast<uint8_t *>(code) + timeout.offset();
   return code;
 }
 
@@ -1841,11 +1875,12 @@ JITX86::JITX86()
   m_pJitEntry = NULL;
 }
 
-bool JITX86::InitializeJIT()
+bool
+JITX86::InitializeJIT()
 {
   g_pCodeCache = KE_CreateCodeCache();
 
-  m_pJitEntry = GenerateEntry(&m_pJitReturn);
+  m_pJitEntry = GenerateEntry(&m_pJitReturn, &m_pJitTimeout);
   if (!m_pJitEntry)
     return false;
 
@@ -1873,6 +1908,11 @@ JITX86::CompileFunction(BaseRuntime *prt, cell_t pcode_offs, int *err)
   JitFunction *fun = cc.emit(err);
   if (!fun)
     return NULL;
+
+  // Grab the lock before linking code in, since the watchdog timer will look
+  // at this list on another thread.
+  ke::AutoLock lock(g_Jit.Mutex());
+
   prt->AddJittedFunction(fun);
   return fun;
 }
@@ -1919,33 +1959,39 @@ JITX86::CreateFakeNative(SPVM_FAKENATIVE_FUNC callback, void *pData)
   return (SPVM_NATIVE_FUNC)LinkCode(masm);
 }
 
-void JITX86::DestroyFakeNative(SPVM_NATIVE_FUNC func)
+void
+JITX86::DestroyFakeNative(SPVM_NATIVE_FUNC func)
 {
   KE_FreeCode(g_pCodeCache, (void *)func);
 }
 
-ICompilation *JITX86::StartCompilation()
+ICompilation *
+JITX86::StartCompilation()
 {
   return new CompData;
 }
 
-ICompilation *JITX86::StartCompilation(BaseRuntime *runtime)
+ICompilation *
+JITX86::StartCompilation(BaseRuntime *runtime)
 {
   return new CompData;
 }
 
-void CompData::Abort()
+void
+CompData::Abort()
 {
   delete this;
 }
 
-void JITX86::FreeContextVars(sp_context_t *ctx)
+void
+JITX86::FreeContextVars(sp_context_t *ctx)
 {
   free(((tracker_t *)(ctx->vm[JITVARS_TRACKER]))->pBase);
   delete (tracker_t *)ctx->vm[JITVARS_TRACKER];
 }
 
-bool CompData::SetOption(const char *key, const char *val)
+bool
+CompData::SetOption(const char *key, const char *val)
 {
   if (strcmp(key, SP_JITCONF_DEBUG) == 0)
     return true;
@@ -1962,7 +2008,8 @@ bool CompData::SetOption(const char *key, const char *val)
   return false;
 }
 
-int JITX86::InvokeFunction(BaseRuntime *runtime, JitFunction *fn, cell_t *result)
+int
+JITX86::InvokeFunction(BaseRuntime *runtime, JitFunction *fn, cell_t *result)
 {
   sp_context_t *ctx = runtime->GetBaseContext()->GetCtx();
 
@@ -1970,18 +2017,75 @@ int JITX86::InvokeFunction(BaseRuntime *runtime, JitFunction *fn, cell_t *result
   ctx->cip = fn->GetPCodeAddress();
 
   JIT_EXECUTE pfn = (JIT_EXECUTE)m_pJitEntry;
+
+  if (level_++ == 0)
+    frame_id_++;
   int err = pfn(ctx, runtime->plugin()->memory, fn->GetEntryAddress());
+  level_--;
 
   *result = ctx->rval;
   return err;
 }
 
-void *JITX86::AllocCode(size_t size)
+void *
+JITX86::AllocCode(size_t size)
 {
   return Knight::KE_AllocCode(g_pCodeCache, size);
 }
 
-void JITX86::FreeCode(void *code)
+void
+JITX86::FreeCode(void *code)
 {
   KE_FreeCode(g_pCodeCache, code);
 }
+
+void
+JITX86::RegisterRuntime(BaseRuntime *rt)
+{
+  mutex_.AssertCurrentThreadOwns();
+  runtimes_.insert(rt);
+}
+
+void
+JITX86::DeregisterRuntime(BaseRuntime *rt)
+{
+  mutex_.AssertCurrentThreadOwns();
+  runtimes_.erase(rt);
+}
+
+void
+JITX86::PatchAllJumpsForTimeout()
+{
+  mutex_.AssertCurrentThreadOwns();
+  for (InlineList<BaseRuntime>::iterator iter = runtimes_.begin(); iter != runtimes_.end(); iter++) {
+    BaseRuntime *rt = *iter;
+    for (size_t i = 0; i < rt->NumJitFunctions(); i++) {
+      JitFunction *fun = rt->GetJitFunction(i);
+      uint8_t *base = reinterpret_cast<uint8_t *>(fun->GetEntryAddress());
+
+      for (size_t j = 0; j < fun->NumLoopEdges(); j++) {
+        const LoopEdge &e = fun->GetLoopEdge(j);
+        int32_t diff = intptr_t(m_pJitTimeout) - intptr_t(base + e.offset);
+        *reinterpret_cast<int32_t *>(base + e.offset - 4) = diff;
+      }
+    }
+  }
+}
+
+void
+JITX86::UnpatchAllJumpsFromTimeout()
+{
+  mutex_.AssertCurrentThreadOwns();
+  for (InlineList<BaseRuntime>::iterator iter = runtimes_.begin(); iter != runtimes_.end(); iter++) {
+    BaseRuntime *rt = *iter;
+    for (size_t i = 0; i < rt->NumJitFunctions(); i++) {
+      JitFunction *fun = rt->GetJitFunction(i);
+      uint8_t *base = reinterpret_cast<uint8_t *>(fun->GetEntryAddress());
+
+      for (size_t j = 0; j < fun->NumLoopEdges(); j++) {
+        const LoopEdge &e = fun->GetLoopEdge(j);
+        *reinterpret_cast<int32_t *>(base + e.offset - 4) = e.disp32;
+      }
+    }
+  }
+}

sourcepawn/jit/x86/jit_x86.h

@@ -1,34 +1,19 @@
-/**
- * vim: set ts=4 sw=4 tw=99 et:
- * =============================================================================
- * SourceMod
- * Copyright (C) 2004-2008 AlliedModders LLC.  All rights reserved.
- * =============================================================================
- *
- * This program is free software; you can redistribute it and/or modify it under
- * the terms of the GNU General Public License, version 3.0, as published by the
- * Free Software Foundation.
- * 
- * This program is distributed in the hope that it will be useful, but WITHOUT
- * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
- * FOR A PARTICULAR PURPOSE.  See the GNU General Public License for more
- * details.
- *
- * You should have received a copy of the GNU General Public License along with
- * this program.  If not, see <http://www.gnu.org/licenses/>.
- *
- * As a special exception, AlliedModders LLC gives you permission to link the
- * code of this program (as well as its derivative works) to "Half-Life 2," the
- * "Source Engine," the "SourcePawn JIT," and any Game MODs that run on software
- * by the Valve Corporation.  You must obey the GNU General Public License in
- * all respects for all other code used.  Additionally, AlliedModders LLC grants
- * this exception to all derivative works.  AlliedModders LLC defines further
- * exceptions, found in LICENSE.txt (as of this writing, version JULY-31-2007),
- * or <http://www.sourcemod.net/license.php>.
- *
- * Version: $Id$
- */
-
+// vim: set ts=8 sts=2 sw=2 tw=99 et:
+//
+// This file is part of SourcePawn.
+// 
+// SourcePawn is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+// 
+// SourcePawn is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+// 
+// You should have received a copy of the GNU General Public License
+// along with SourcePawn.  If not, see <http://www.gnu.org/licenses/>.
 #ifndef _INCLUDE_SOURCEPAWN_JIT_X86_H_
 #define _INCLUDE_SOURCEPAWN_JIT_X86_H_
 
@@ -42,6 +27,7 @@
 #include "sp_vm_basecontext.h"
 #include "jit_function.h"
 #include "opcodes.h"
+#include <ke_thread_utils.h>
 
 using namespace SourcePawn;
 
@@ -155,6 +141,7 @@ class Compiler
   cell_t *cip_;
   cell_t *code_end_;
   Label *jump_map_;
+  ke::Vector<size_t> backward_jumps_;
 
   // Errors
   Label error_bounds_;
@@ -188,16 +175,36 @@ class JITX86
   ICompilation *ApplyOptions(ICompilation *_IN, ICompilation *_OUT);
   int InvokeFunction(BaseRuntime *runtime, JitFunction *fn, cell_t *result);
 
+  void RegisterRuntime(BaseRuntime *rt);
+  void DeregisterRuntime(BaseRuntime *rt);
+  void PatchAllJumpsForTimeout();
+  void UnpatchAllJumpsFromTimeout();
+  
  public:
   ExternalAddress GetUniversalReturn() {
-      return ExternalAddress(m_pJitReturn);
+    return ExternalAddress(m_pJitReturn);
   }
   void *AllocCode(size_t size);
   void FreeCode(void *code);
 
+  uintptr_t FrameId() const {
+    return frame_id_;
+  }
+  bool RunningCode() const {
+    return level_ != 0;
+  }
+  ke::Mutex *Mutex() {
+    return &mutex_;
+  }
+
  private:
   void *m_pJitEntry;         /* Entry function */
   void *m_pJitReturn;        /* Universal return address */
+  void *m_pJitTimeout;       /* Universal timeout address */
+  InlineList<BaseRuntime> runtimes_;
+  uintptr_t frame_id_;
+  uintptr_t level_;
+  ke::Mutex mutex_;
 };
 
 const Register pri = eax;
@@ -211,3 +218,4 @@ extern Knight::KeCodeCache *g_pCodeCache;
 extern JITX86 g_Jit;
 
 #endif //_INCLUDE_SOURCEPAWN_JIT_X86_H_
+