Fix #1

README.md

@@ -2,7 +2,6 @@
 
 Add `await` and `async` keyword in Crystal.
 
-
 ## Usage
 
 - Call `async` on any method or block to create a MiniFuture
@@ -12,6 +11,18 @@ Add `await` and `async` keyword in Crystal.
 Can improve drastically application which relay on blocking IO like web API
 or file writing.
 
+### await(timeout, future)
+
+```crystal
+f = async check_website
+
+begin
+  await 5.seconds, f
+rescue MiniFuture::TimeoutException
+  # rescue from timeout
+end
+```
+
 ### async_lp
 
 By default, async call the newly created fiber just after creation.

spec/await_async_spec.cr

@@ -61,5 +61,27 @@ module AwaitAsyncSpec
       await f
       await f
     end
+
+    it "can await with timeout" do
+      x = async { sleep 0.5 }
+
+      expect_raises MiniFuture::TimeoutException do
+        await 0.2.seconds, x
+      end
+
+      x = async { sleep 0.1 }
+      await 0.2.seconds, x # < Should not raise exception
+    end
+  end
+
+  it "can await with timeout (array)" do
+    x = 5.times.map { |arr| async { sleep 0.5 } }.to_a
+
+    expect_raises MiniFuture::TimeoutException do
+      await 0.2.seconds, x
+    end
+
+    x = 5.times.map { async { sleep 0.1 } }.to_a
+    await 0.2.seconds, x # < Should not raise exception
   end
 end

src/await_async.cr

@@ -1,5 +1,8 @@
 # Lightweight Future structure.
 class MiniFuture(T)
+  class TimeoutException < Exception
+  end
+
   @status = :running
   @channel = Channel(T).new(1)
   @error : Exception? = nil
@@ -31,10 +34,26 @@ class MiniFuture(T)
     @status == :running
   end
 
-  def await
+  def await(timeout : Time::Span? = nil)
     if @status != :flushed
-      @status = :flushed
-      @output = @channel.receive
+      if timeout
+        timeout_channel = Channel(Nil).new
+
+        spawn do
+          sleep timeout.not_nil!
+          timeout_channel.send nil unless @status == :flushed
+        end
+
+        select
+        when timeout_channel.receive
+          raise TimeoutException.new
+        when @output = @channel.receive
+          @status = :flushed
+        end
+      else
+        @status = :flushed
+        @output = @channel.receive
+      end
     end
 
     if e = @error
@@ -53,8 +72,16 @@ def await(future : MiniFuture)
   future.await
 end
 
-def await(futures : Iterator(MiniFuture(T))) forall T
-  futures.map { |x| await x }
+def await(timeout : Time::Span, future : MiniFuture)
+  future.await(timeout)
+end
+
+def await(futures : Iterator(MiniFuture(T)) | Array(MiniFuture(T))) forall T
+  futures.map(&.await)
+end
+
+def await(timeout : Time::Span, futures : Iterator(MiniFuture(T)) | Array(MiniFuture(T))) forall T
+  futures.map(&.await(timeout))
 end
 
 macro async(method)