From 24035602f256bb421e4347b90dce9951b434f3cd Mon Sep 17 00:00:00 2001
From: Wojtek Mach <wojtek@wojtekmach.pl>
Date: Wed, 18 Jan 2023 14:34:46 +0100
Subject: [PATCH] Update demo

---
 elixirkit/demo/lib/demo.ex | 12 ++++++++----
 1 file changed, 8 insertions(+), 4 deletions(-)

diff --git a/elixirkit/demo/lib/demo.ex b/elixirkit/demo/lib/demo.ex
index 54dd0324d..7ed3e69ef 100644
--- a/elixirkit/demo/lib/demo.ex
+++ b/elixirkit/demo/lib/demo.ex
@@ -24,7 +24,10 @@ defmodule Demo.Server do
   @impl true
   def init(_) do
     Process.flag(:trap_exit, true)
-    {:ok, server_pid} = ElixirKit.start()
+
+    {:ok, pid} = ElixirKit.start()
+    ref = Process.monitor(pid)
+
     log("init")
 
     Task.start(fn ->
@@ -36,7 +39,7 @@ defmodule Demo.Server do
       System.stop()
     end)
 
-    {:ok, %{server_pid: server_pid}}
+    {:ok, %{ref: ref}}
   end
 
   @impl true
@@ -46,8 +49,9 @@ defmodule Demo.Server do
   end
 
   @impl true
-  def handle_info({:EXIT, pid, :shutdown}, state) when pid == state.server_pid do
-    {:noreply, state}
+  def handle_info({:DOWN, ref, :process, _, :shutdown}, state) when ref == state.ref do
+    System.stop()
+    {:stop, :shutdown, state}
   end
 
   @impl true