feat: Correctly handle kill events

core/homescript/manager.go

@@ -42,11 +42,14 @@ type Manager struct {
 }
 
 type Job struct {
-	Username  string
-	JobId     uint64
-	HmsId     *string
-	Initiator HomescriptInitiator
-	CancelCtx context.CancelFunc
+	Username        string
+	JobId           uint64
+	HmsId           *string
+	Initiator       HomescriptInitiator
+	CancelCtx       context.CancelFunc
+	Interpreter     interpreter.Interpreter
+	EntryModuleName string
+	SupportsKill    bool
 }
 
 // For external usage (can be marshaled)
@@ -113,15 +116,21 @@ func (m *Manager) PushJob(
 	initiator HomescriptInitiator,
 	cancelCtxFunc context.CancelFunc,
 	hmsId *string,
+	interpreter *interpreter.Interpreter,
+	entryModuleName string,
+	supportsKill bool,
 ) uint64 {
 	m.Lock.Lock()
 	id := uint64(len(m.Jobs))
 	m.Jobs = append(m.Jobs, Job{
-		Username:  username,
-		JobId:     id,
-		HmsId:     hmsId,
-		Initiator: initiator,
-		CancelCtx: cancelCtxFunc,
+		Username:        username,
+		JobId:           id,
+		HmsId:           hmsId,
+		Initiator:       initiator,
+		CancelCtx:       cancelCtxFunc,
+		Interpreter:     *interpreter,
+		EntryModuleName: entryModuleName,
+		SupportsKill:    supportsKill,
 	})
 	m.Lock.Unlock()
 	return id
@@ -248,42 +257,48 @@ func (m *Manager) Run(
 ) (HmsRes, error) {
 	// TODO: handle arguments
 
-	id := m.PushJob(username, initiator, cancelCtxFunc, filename)
-	defer m.removeJob(id)
-
-	internalFilename := fmt.Sprintf("live@%d", id) // TODO: the @ symbol cannot be used in IDs?
+	// TODO: the @ symbol cannot be used in IDs?
+	// FIX: implement this uniqueness properly
+	entryModuleName := fmt.Sprintf("live@%d", time.Now().Nanosecond())
 	if filename != nil {
-		internalFilename = *filename
+		entryModuleName = *filename
 	}
 
-	modules, res, err := m.Analyze(username, internalFilename, code)
+	modules, res, err := m.Analyze(username, entryModuleName, code)
 	if err != nil {
 		return HmsRes{}, err
 	}
 	if !res.Success {
 		return res, nil
 	}
 
-	// send the id to the id channel (only if it exists)
-	if idChan != nil {
-		*idChan <- id
-	}
+	log.Debug(fmt.Sprintf("Homescript '%s' of user '%s' is executing...", entryModuleName, username))
 
-	log.Debug(fmt.Sprintf("Homescript '%s' of user '%s' is executing...", internalFilename, username))
-	if i := homescript.Run(
+	interpreter := interpreter.NewInterpreter(
 		CALL_STACK_LIMIT_SIZE,
-		modules,
-		internalFilename,
 		newInterpreterExecutor(
 			username,
 			outputWriter,
 			args,
 			automationContext,
 			cancelCtxFunc,
 		),
+		modules,
 		interpreterScopeAdditions(),
 		&cancelCtx,
-	); i != nil {
+	)
+
+	supportsKill := modules[entryModuleName].SupportsEvent("kill")
+
+	id := m.PushJob(username, initiator, cancelCtxFunc, filename, &interpreter, entryModuleName, supportsKill)
+	defer m.removeJob(id)
+
+	// send the id to the id channel (only if it exists)
+	if idChan != nil {
+		*idChan <- id
+	}
+
+	if i := interpreter.Execute(entryModuleName); i != nil {
 		span := errors.Span{}
 
 		i := *i
@@ -331,14 +346,14 @@ func (m *Manager) Run(
 		}
 
 		if isErr {
-			fileContentsTemp, err := resolveFileContentsOfErrors(username, internalFilename, code, errors)
+			fileContentsTemp, err := resolveFileContentsOfErrors(username, entryModuleName, code, errors)
 			if err != nil {
 				return HmsRes{}, err
 			}
 
 			fileContents = fileContentsTemp
 
-			log.Debug(fmt.Sprintf("Homescript '%s' of user '%s' failed: %s", internalFilename, username, errors[0]))
+			log.Debug(fmt.Sprintf("Homescript '%s' of user '%s' failed: %s", entryModuleName, username, errors[0]))
 		}
 
 		return HmsRes{
@@ -348,7 +363,7 @@ func (m *Manager) Run(
 		}, nil
 	}
 
-	log.Debug(fmt.Sprintf("Homescript '%s' of user '%s' executed successfully", internalFilename, username))
+	log.Debug(fmt.Sprintf("Homescript '%s' of user '%s' executed successfully", entryModuleName, username))
 
 	return HmsRes{Success: true, Errors: make([]HmsError, 0), FileContents: make(map[string]string)}, nil
 }
@@ -426,9 +441,7 @@ func (m *Manager) Kill(jobId uint64) bool {
 	defer m.Lock.Unlock()
 	for _, job := range m.Jobs {
 		if job.JobId == jobId {
-			log.Trace("Dispatching sigTerm to HMS interpreter channel")
-			job.CancelCtx()
-			log.Trace("Successfully dispatched sigTerm to HMS interpreter channel")
+			m.killJob(job)
 			return true
 		}
 	}
@@ -446,16 +459,20 @@ func (m *Manager) KillAllId(hmsId string) (count uint64, success bool) {
 		}
 
 		// Only standalone scripts may be terminated (callstack validation) | TODO: implement this
-		log.Trace("Dispatching sigTerm to HMS interpreter channel")
-		job.CancelCtx()
-		log.Trace("Successfully dispatched sigTerm to HMS interpreter channel")
+		m.killJob(job)
 
 		success = true
 		count++
 	}
 	return count, success
 }
 
+func (m *Manager) killJob(job Job) {
+	log.Trace("Dispatching sigTerm to HMS interpreter channel")
+	job.CancelCtx()
+	log.Trace("Successfully dispatched sigTerm to HMS interpreter channel")
+}
+
 // Can be used to access the manager's jobs from the outside in a safe manner
 func (m *Manager) GetJobList() []Job {
 	m.Lock.RLock()

server/api/homescriptAsync.go

@@ -123,16 +123,19 @@ func RunHomescriptByIDAsync(w http.ResponseWriter, r *http.Request) {
 
 	// Start running the code
 	res := make(chan homescript.HmsRes)
-	ctx, cancel := context.WithCancel(context.Background())
 
-	go func(writer io.Writer, results *chan homescript.HmsRes, ctx context.Context, cancel context.CancelFunc) {
+	idChan := make(chan uint64)
+
+	go func(writer io.Writer, results *chan homescript.HmsRes, idChan *chan uint64) {
+		ctx, cancel := context.WithCancel(context.Background())
+
 		res, err := homescript.HmsManager.RunById(
 			request.Payload,
 			username,
 			homescript.InitiatorAPI,
 			ctx,
 			cancel,
-			nil,
+			idChan,
 			args,
 			outWriter,
 			nil,
@@ -152,7 +155,9 @@ func RunHomescriptByIDAsync(w http.ResponseWriter, r *http.Request) {
 		outWriter.Close()
 
 		*results <- res
-	}(outWriter, &res, ctx, cancel)
+	}(outWriter, &res, &idChan)
+
+	jobId := <-idChan
 
 	go func() {
 		// Check if the script should be killed
@@ -185,7 +190,8 @@ func RunHomescriptByIDAsync(w http.ResponseWriter, r *http.Request) {
 		}
 		// Kill the Homescript
 		log.Trace("Killing script via Websocket")
-		cancel()
+		// cancel()
+		homescript.HmsManager.Kill(jobId)
 		log.Trace("Killed script via Websocket")
 	}()