change shutdown

cmd/aws-lambda-rie/handlers.go

@@ -72,7 +72,7 @@ func printEndReports(invokeId string, initDuration string, memorySize string, in
 		invokeId, invokeDuration, math.Ceil(invokeDuration), memorySize, memorySize)
 }
 
-func InvokeHandler(w http.ResponseWriter, r *http.Request, sandbox Sandbox, bs interop.Bootstrap, done chan<- struct{}) {
+func InvokeHandler(w http.ResponseWriter, r *http.Request, sandbox Sandbox, bs interop.Bootstrap, doneCallback func()) {
 	log.Debugf("invoke: -> %s %s %v", r.Method, r.URL, r.Header)
 	bodyBytes, err := ioutil.ReadAll(r.Body)
 	if err != nil {
@@ -191,7 +191,7 @@ func InvokeHandler(w http.ResponseWriter, r *http.Request, sandbox Sandbox, bs i
 	}
 	w.Write(invokeResp.Body)
 
-	close(done) // Signal that the InvokeHandler has completed its execution
+	doneCallback() // Callback at end of InvokeHandler
 }
 
 func InitHandler(sandbox Sandbox, functionVersion string, timeout int64, bs interop.Bootstrap) (time.Time, time.Time) {

cmd/aws-lambda-rie/http.go

@@ -5,7 +5,6 @@ package main
 
 import (
 	"os"
-	"time"
 	"strconv"
 	"net/http"
 
@@ -35,24 +34,20 @@ func startHTTPServer(ipport string, sandbox *rapidcore.SandboxBuilder, bs intero
 
 	// Pass a channel
 	http.HandleFunc("/2015-03-31/functions/function/invocations", func(w http.ResponseWriter, r *http.Request) {
-		done := make(chan struct{}) // Channel to signal when the response has terminated
-
-		InvokeHandler(w, r, sandbox.LambdaInvokeAPI(), bs, done)
-
-		// Shutdown the server if we've reached the maximum number of invocations
-		maxInvocations--
-		if maxInvocations == 0 {
-			<-done // Wait until the response has terminated
-			close(shutdownChan)
-		}
+		InvokeHandler(w, r, sandbox.LambdaInvokeAPI(), bs, func(){
+			// Shutdown the server if the maximum number of invocations is reached
+			maxInvocations--
+			if maxInvocations == 0 {
+				close(shutdownChan)
+			}
+		})
 	})
 
 
 	// go routine to handle server shutdown (main thread waits)
 	go func() {
 		<-shutdownChan
 		log.Printf("Maximum invocations (%s) reached. Shutting down the server.", maxInvocationsStr)
-		time.Sleep(60 * time.Second) // Grace period for client to receive the response
 		if err := srv.Shutdown(nil); err != nil {
 			log.Panic(err)
 		}