1#!/usr/bin/env perl
2
3use strict;
4use warnings;
5use Test::More;
6use FindBin qw($Bin);
7use lib "$Bin/lib";
8use Carp qw(croak);
9use MemcachedTest;
10use IO::Select;
11use IO::Socket qw(AF_INET SOCK_STREAM);
12
13# TODO: possibly... set env var to a generated temp filename before starting
14# the server so we can pass that in?
15my $configfile = "/tmp/proxyustats.lua";
16
17if (!supports_proxy()) {
18    plan skip_all => 'proxy not enabled';
19    exit 0;
20}
21
22# Set up some server sockets.
23sub mock_server {
24    my $port = shift;
25    my $srv = IO::Socket->new(
26        Domain => AF_INET,
27        Type => SOCK_STREAM,
28        Proto => 'tcp',
29        LocalHost => '127.0.0.1',
30        LocalPort => $port,
31        ReusePort => 1,
32        Listen => 5) || die "IO::Socket: $@";
33    return $srv;
34}
35
36sub accept_backend {
37    my $srv = shift;
38    my $be = $srv->accept();
39    $be->autoflush(1);
40    ok(defined $be, "mock backend created");
41    like(<$be>, qr/version/, "received version command");
42    print $be "VERSION 1.0.0-mock\r\n";
43
44    return $be;
45}
46
47# Put a version command down the pipe to ensure the socket is clear.
48# client version commands skip the proxy code
49sub check_version {
50    my $ps = shift;
51    print $ps "version\r\n";
52    like(<$ps>, qr/VERSION /, "version received");
53}
54
55# A single line containing prefix, start ustat and end ustat index.
56sub write_config {
57    my $cmd = shift;
58    open(my $fh, "> $configfile") or die "Couldn't overwrite $configfile: $!";
59    print $fh $cmd;
60    close($fh);
61}
62
63sub wait_reload {
64    my $w = shift;
65    like(<$w>, qr/ts=(\S+) gid=\d+ type=proxy_conf status=start/, "reload started");
66    like(<$w>, qr/ts=(\S+) gid=\d+ type=proxy_conf status=done/, "reload completed");
67}
68
69my ($p_srv, $ps, $watcher);
70sub restart_memcached {
71    if ($p_srv) {
72        $p_srv->stop();
73    }
74    write_config('return "a 1 0"');
75    $p_srv = new_memcached('-o proxy_config=./t/proxyustats.lua -t 1');
76    $ps = $p_srv->sock;
77    $ps->autoflush(1);
78
79    $watcher = $p_srv->new_sock;
80    print $watcher "watch proxyevents\n";
81    is(<$watcher>, "OK\r\n", "watcher enabled");
82}
83
84diag "testing failure to start";
85write_config("invalid");
86eval {
87    $p_srv = new_memcached('-o proxy_config=./t/proxyustats.lua -t 1');
88};
89ok($@ && $@ =~ m/Failed to connect/, "server successfully not started");
90
91restart_memcached();
92
93subtest 'succeeded to allocate 1024 ustats' => sub {
94    my $pfx = "a";
95    my $first = 1;
96    my $last = 1024;
97    write_config('return "' . $pfx . ' ' . $first . ' ' . $last . '"');
98    $p_srv->reload();
99    wait_reload($watcher);
100
101    my $stats = mem_stats($ps, 'proxy');
102    for my $i ($first..$last) {
103        my $ustat = "user_" . $pfx . $i;
104        is($stats->{$ustat}, 0, $ustat . " found");
105    }
106};
107
108subtest 'failed to allocate 1025 ustats' => sub {
109    my $pfx = "a";
110    my $first = 1025;
111    my $last = 1025;
112    write_config('return "' . $pfx . ' ' . $first . ' ' . $last . '"');
113    $p_srv->reload();
114    unlike(<$watcher>, qr/ts=(\S+) gid=\d+ type=proxy_conf status=start/, "reload not started");
115    unlike(<$watcher>, qr/ts=(\S+) gid=\d+ type=proxy_conf status=done/, "reload not completed");
116
117    restart_memcached();
118};
119
120subtest 'failed to allocate ustat at index 0' => sub {
121    my $pfx = "a";
122    my $first = 0;
123    my $last = 0;
124    write_config('return "' . $pfx . ' ' . $first . ' ' . $last . '"');
125    $p_srv->reload();
126    unlike(<$watcher>, qr/ts=(\S+) gid=\d+ type=proxy_conf status=start/, "reload not started");
127    unlike(<$watcher>, qr/ts=(\S+) gid=\d+ type=proxy_conf status=done/, "reload not completed");
128
129    restart_memcached();
130};
131
132
133subtest 'succeeded to allocate ustat at 1024 only' => sub {
134    # restart memcached to clear any ustats.
135    restart_memcached();
136
137    my $pfx = "a";
138    my $first = 1024;
139    my $last = 1024;
140    write_config('return "' . $pfx . ' ' . $first . ' ' . $last . '"');
141    $p_srv->reload();
142    wait_reload($watcher);
143
144    my $stats = mem_stats($ps, 'proxy');
145    while (my ($ustat, $value) = each %{$stats}) {
146        if (index($ustat, "user_") == 0) {
147            is($ustat, "user_a1024", "user_a1024 found");
148        }
149    }
150};
151
152subtest 'ustats incr/decr and perseverance over reload' => sub {
153    # restart memcached to clear any ustats.
154    restart_memcached();
155
156    my $pfx = "a";
157    my $first = 1;
158    my $last = 3;
159    write_config('return "' . $pfx . ' ' . $first . ' ' . $last . '"');
160    $p_srv->reload();
161    wait_reload($watcher);
162
163    print $ps "mg 1\r\n";
164    is(scalar <$ps>, "HD\r\n", "mg 1 hit");
165    print $ps "mg 2\r\n";
166    is(scalar <$ps>, "HD\r\n", "mg 2 hit");
167    print $ps "mg 2\r\n";
168    is(scalar <$ps>, "HD\r\n", "mg 2 hit");
169
170    my $stats = mem_stats($ps, 'proxy');
171    is($stats->{user_a1}, 1, "user_a1 is 1");
172    is($stats->{user_a2}, 4, "user_a2 is 4");
173    is($stats->{user_a3}, 0, "user_a3 is 0");
174
175    $pfx = "b";
176    $first = 2;
177    $last = 4;
178    write_config('return "' . $pfx . ' ' . $first . ' ' . $last . '"');
179    $p_srv->reload();
180    wait_reload($watcher);
181
182    # add 2 at idx 2.
183    print $ps "mg 2\r\n";
184    is(scalar <$ps>, "HD\r\n", "mg 2 hit");
185
186    $stats = mem_stats($ps, 'proxy');
187    # carried over since name did not change
188    is($stats->{user_a1}, 1, "user_a1 is 1");
189    # index 2 changed names, so should be reset.
190    is($stats->{user_b2}, 2, "user_b2 is 2 instead of 4");
191    is($stats->{user_b3}, 0, "user_b3 is 0");
192    is($stats->{user_b4}, 0, "user_b4 is 0");
193};
194
195subtest 'failed to allocate ustat longer than 128 chars' => sub {
196    my $pfx = '*' x 128;
197    my $first = 1;
198    my $last = 1;
199    write_config('return "' . $pfx . ' ' . $first . ' ' . $last . '"');
200    $p_srv->reload();
201    unlike(<$watcher>, qr/ts=(\S+) gid=\d+ type=proxy_conf status=start/, "reload not started");
202    unlike(<$watcher>, qr/ts=(\S+) gid=\d+ type=proxy_conf status=done/, "reload not completed");
203
204    restart_memcached();
205};
206
207subtest 'negative ustat value underflow' => sub {
208    # restart memcached to clear any ustats.
209    restart_memcached();
210
211    my $pfx = "a";
212    my $first = 1;
213    my $last = 1;
214    write_config('return "' . $pfx . ' ' . $first . ' ' . $last . '"');
215    $p_srv->reload();
216    wait_reload($watcher);
217
218    print $ps "mg -1\r\n";
219    is(scalar <$ps>, "HD\r\n", "mg -1 hit");
220
221    my $stats = mem_stats($ps, 'proxy');
222    isnt($stats->{user_a1}, -1, "user_a1 is not -1");
223};
224
225subtest 'mask off stat after reload' => sub {
226    restart_memcached();
227
228    write_config('return "mask1 1 1"');
229    $p_srv->reload();
230    wait_reload($watcher);
231
232    my $s1 = mem_stats($ps, 'proxy');
233    is($s1->{"user_maska"}, 0, "maska is 0");
234    is($s1->{"user_maskb"}, 0, "maskb is 0");
235
236    write_config('return "mask2 1 1"');
237    $p_srv->reload();
238    wait_reload($watcher);
239
240    my $s2 = mem_stats($ps, 'proxy');
241
242    is($s2->{"user_maska"}, undef, "maska is gone");
243    is($s2->{"user_maskb"}, 0, "maskb is 0");
244};
245
246done_testing();
247
248END {
249    unlink $configfile;
250}
251