From 1c6fe289573aaae41db68d3d8046156e8d3f34eb Mon Sep 17 00:00:00 2001
From: Greg Furman <gregfurman99@gmail.com>
Date: Thu, 26 Jun 2025 18:38:07 +0200
Subject: [PATCH] feat(init): Handle expiring init

---
 internal/server/interop.go | 30 ++++++++++++++++++++++++++++++
 1 file changed, 30 insertions(+)

diff --git a/internal/server/interop.go b/internal/server/interop.go
index 9bb8f84..6a555fa 100644
--- a/internal/server/interop.go
+++ b/internal/server/interop.go
@@ -6,6 +6,7 @@ import (
 	"errors"
 	"fmt"
 	"net/http"
+	"time"
 
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/localstack/lambda-runtime-init/internal/localstack"
@@ -31,6 +32,35 @@ func NewInteropServer(ls *localstack.LocalStackClient) *LocalStackInteropsServer
 	}
 }
 
+func (c *LocalStackInteropsServer) Init(initRequest *interop.Init, timeoutMs int64) error {
+	// This allows us to properly timeout when an INIT request -- which is unimplemented in the upstream.
+
+	initStart := metering.Monotime()
+
+	initDone := make(chan error, 1)
+	go func() {
+		initDone <- c.Server.Init(initRequest, timeoutMs)
+	}()
+
+	var err error
+	select {
+	case err = <-initDone:
+	case <-time.After(time.Duration(timeoutMs) * time.Millisecond):
+		if _, resetErr := c.Server.Reset("timeout", 2000); resetErr != nil {
+			log.WithError(resetErr).Error("Failed to reset after init timeout")
+		}
+		err = errors.New("timeout")
+	}
+
+	initDuration := float64(metering.Monotime()-initStart) / float64(time.Millisecond)
+
+	if err != nil {
+		log.WithError(err).WithField("duration", initDuration).Error("Init failed")
+	}
+
+	return err
+}
+
 func (c *LocalStackInteropsServer) Execute(ctx context.Context, responseWriter http.ResponseWriter, invoke *interop.Invoke) error {
 	ctx, cancel := context.WithTimeout(context.Background(), c.Server.GetInvokeTimeout())
 	defer cancel()