Support psgix.cleanup (and psgix.cleanup.handlers)

lib/Starman.pm

@@ -67,7 +67,11 @@ children (workers) is less than 3.0MB.
 
 =item PSGI compatible
 
-Can run any PSGI applications and frameworks
+Can run any PSGI applications and frameworks.
+
+It also supports the C<psgix.cleanup> and C<psgix.harakiri> as described in
+L<PSGI::Extensions>. Note that use of the cleanup handler will turn off HTTP
+keep-alive.
 
 =item HTTP/1.1 support
 

lib/Starman/Server.pm

@@ -234,6 +234,8 @@ sub process_request {
             'psgix.io'          => $conn,
             'psgix.input.buffered' => Plack::Util::TRUE,
             'psgix.harakiri' => Plack::Util::TRUE,
+            'psgix.cleanup' => Plack::Util::TRUE,
+            'psgix.cleanup.handlers' => [],
         };
 
         # Parse headers
@@ -467,11 +469,23 @@ sub _prepare_env {
 sub _finalize_response {
     my($self, $env, $res) = @_;
 
+    # We also check for harakiri in post_client_connection_hook()
+    # after the cleanup handlers have run.
     if ($env->{'psgix.harakiri.commit'}) {
         $self->{client}->{keepalive} = 0;
         $self->{client}->{harakiri} = 1;
     }
 
+    if (@{$env->{'psgix.cleanup.handlers'}}) {
+        $self->{client}->{env_and_cleanup_handlers} = {
+            env => $env,
+            cleanup_handlers => $env->{'psgix.cleanup.handlers'}
+        };
+        # If keepalive stayed on, we could serve one or more requests
+        # before reaching post_client_connection_hook()
+        $self->{client}->{keepalive} = 0;
+    }
+
     my $protocol = $env->{SERVER_PROTOCOL};
     my $status   = $res->[0];
     my $message  = status_message($status);
@@ -581,6 +595,21 @@ sub _syswrite {
 
 sub post_client_connection_hook {
     my $self = shift;
+
+   if ($self->{client}->{env_and_cleanup_handlers}) {
+        my ($env, $cleanup_handlers) = @{$self->{client}->{env_and_cleanup_handlers}}{qw(env cleanup_handlers)};
+
+        for my $cleanup_handler (@$cleanup_handlers) {
+            $cleanup_handler->($env);
+        }
+
+        # We also check for harakiri in _finalize_response()
+        # after the request has been completed.
+        if ($env->{'psgix.harakiri.commit'}) {
+            $self->{client}->{harakiri} = 1;
+        }
+    }
+
     if ($self->{client}->{harakiri}) {
         exit;
     }

t/harakiri.t

@@ -39,4 +39,23 @@ test_psgi
         is keys(%seen_pid), 23, 'In Harakiri mode, each pid only used once';
     };
 
+test_psgi
+    app => sub {
+        my $env = shift;
+        push @{$env->{'psgix.cleanup.handlers'}} => sub {
+            my ($cleanup_env) = @_;
+            $cleanup_env->{'psgix.harakiri.commit'} = 1;
+        };
+        return [ 200, [ 'Content-Type' => 'text/plain' ], [$$] ];
+    },
+    client => sub {
+        my %seen_pid;
+        my $cb = shift;
+        for (1..23) {
+            my $res = $cb->(GET "/");
+            $seen_pid{$res->content}++;
+        }
+        is keys(%seen_pid), 23, 'In Harakiri mode triggered by cleanup, each pid only used once';
+    };
+
 done_testing;

t/psgix_cleanup.t

@@ -0,0 +1,66 @@
+use strict;
+use warnings;
+
+use HTTP::Request::Common;
+use Plack::Test;
+use Test::More;
+
+$Plack::Test::Impl = 'Server';
+$ENV{PLACK_SERVER} = 'Starman';
+
+test_psgi
+    app => sub {
+        my $env = shift;
+        return [ 200, [ 'Content-Type' => 'text/plain' ], [$env->{'psgix.cleanup'}] ];
+    },
+    client => sub {
+        my $cb = shift;
+        my $res = $cb->(GET "/");
+        ok($res->content, "We set psgix.cleanup");
+    };
+
+test_psgi
+    app => sub {
+        my $env = shift;
+        return [ 200, [ 'Content-Type' => 'text/plain' ], [ref($env->{'psgix.cleanup.handlers'})] ];
+    },
+    client => sub {
+        my $cb = shift;
+        my $res = $cb->(GET "/");
+        cmp_ok($res->content, "eq", "ARRAY", "psgix.cleanup.handlers is an array");
+    };
+
+test_psgi
+    app => sub {
+        my $env = shift;
+        return [ 200, [ 'Content-Type' => 'text/plain' ], [join "", @{$env->{'psgix.cleanup.handlers'}}] ];
+    },
+    client => sub {
+        my $cb = shift;
+        my $res = $cb->(GET "/");
+        cmp_ok($res->content, "eq", "", "..which is empty by default");
+    };
+
+my $content = "NO_CLEANUP";
+test_psgi
+    app => sub {
+        my $env = shift;
+        push @{$env->{'psgix.cleanup.handlers'}} => sub { $content .= "XXX" };
+        push @{$env->{'psgix.cleanup.handlers'}} => sub { $content .= "YYY" };
+        push @{$env->{'psgix.cleanup.handlers'}} => sub { $content .= "ZZZ" };
+        return [ 200, [ 'Content-Type' => 'text/plain' ], [ $content ] ];
+    },
+    client => sub {
+        my $cb = shift;
+        my $res = $cb->(GET "/");
+        cmp_ok($res->content, "eq", "NO_CLEANUP", "By the time we run the cleanup handler we've already returned a response");
+
+        my $responses;
+        for (1..10) {
+            my $res = $cb->(GET "/");
+            $responses .= $res->content;
+        }
+        like($responses, qr/$_/, "The response contains '$_' indicating the cleanup handlers were run") for qw(XXX YYY ZZZ);
+    };
+
+done_testing;