From 5f0fcb01dd8b992412a91a4015356e9d4fc27dab Mon Sep 17 00:00:00 2001
From: syumai <syumai@gmail.com>
Date: Fri, 8 Nov 2024 01:46:52 +0900
Subject: [PATCH] split workers.Ready and workers.WaitForCompletion from
 workers.Serve

---
 cloudflare/cron/scheduler.go |  4 +++-
 handler.go                   | 12 ++++++++++++
 handler_js.go                | 24 ++++++++++++++++++++----
 3 files changed, 35 insertions(+), 5 deletions(-)

diff --git a/cloudflare/cron/scheduler.go b/cloudflare/cron/scheduler.go
index 1f97881..0c0bd0d 100644
--- a/cloudflare/cron/scheduler.go
+++ b/cloudflare/cron/scheduler.go
@@ -58,4 +58,6 @@ func ScheduleTask(task Task) {
 
 // ScheduleTaskNonBlock sets the Task to be executed but does not signal readiness or block
 // indefinitely. The non-blocking form is meant to be used in conjunction with [workers.Serve].
-func ScheduleTaskNonBlock(task Task) { scheduledTask = task }
+func ScheduleTaskNonBlock(task Task) {
+	scheduledTask = task
+}
diff --git a/handler.go b/handler.go
index edc3f5b..2e195eb 100644
--- a/handler.go
+++ b/handler.go
@@ -25,3 +25,15 @@ func Serve(handler http.Handler) {
 	fmt.Fprintln(os.Stderr, "warn: this server is currently running in non-JS mode. to enable JS-related features, please use the make command in the syumai/workers template.")
 	http.ListenAndServe(addr, handler)
 }
+
+func ServeNonBlock(http.Handler) {
+	panic("ServeNonBlock is not supported in non-JS environments")
+}
+
+func Ready() {
+	panic("Ready is not supported in non-JS environments")
+}
+
+func WaitForCompletion() {
+	panic("WaitForCompletion is not supported in non-JS environments")
+}
diff --git a/handler_js.go b/handler_js.go
index ed026a3..9faa52f 100644
--- a/handler_js.go
+++ b/handler_js.go
@@ -84,16 +84,32 @@ func handleRequest(reqObj js.Value) (js.Value, error) {
 	return w.ToJSResponse(), nil
 }
 
-//go:wasmimport workers ready
-func ready()
-
-// Server serves http.Handler on a JS runtime.
+// Serve serves http.Handler on a JS runtime.
 // if the given handler is nil, http.DefaultServeMux will be used.
 func Serve(handler http.Handler) {
+	ServeNonBlock(handler)
+	Ready()
+	WaitForCompletion()
+}
+
+// ServeNonBlock sets the http.Handler to be served but does not signal readiness or block
+// indefinitely. The non-blocking form is meant to be used in conjunction with Ready and WaitForCompletion.
+func ServeNonBlock(handler http.Handler) {
 	if handler == nil {
 		handler = http.DefaultServeMux
 	}
 	httpHandler = handler
+}
+
+//go:wasmimport workers ready
+func ready()
+
+// Ready must be called after all setups of the Go side's handlers are done.
+func Ready() {
 	ready()
+}
+
+// WaitForCompletion blocks until the handler set by ServeNonBlock is completed.
+func WaitForCompletion() {
 	<-closeCh
 }