lib/PublicInbox/SharedKV.pm

   1 # Copyright (C) 2020-2021 all contributors <meta@public-inbox.org>
   2 # License: AGPL-3.0+ <https://www.gnu.org/licenses/agpl-3.0.txt>
   3
   4 # fork()-friendly key-value store.  Will be used for making
   5 # augmenting Maildirs and mboxes less expensive, maybe.
   6 # We use flock(2) to avoid SQLite lock problems (busy timeouts, backoff)
   7 package PublicInbox::SharedKV;
   8 use strict;
   9 use v5.10.1;
  10 use parent qw(PublicInbox::Lock);
  11 use File::Temp qw(tempdir);
  12 use DBI ();
  13 use PublicInbox::Spawn;
  14 use File::Path qw(rmtree make_path);
  15
  16 sub dbh {
  17         my ($self, $lock) = @_;
  18         $self->{dbh} //= do {
  19                 my $f = $self->{filename};
  20                 $lock //= $self->lock_for_scope_fast;
  21                 my $dbh = DBI->connect("dbi:SQLite:dbname=$f", '', '', {
  22                         AutoCommit => 1,
  23                         RaiseError => 1,
  24                         PrintError => 0,
  25                         sqlite_use_immediate_transaction => 1,
  26                         # no sqlite_unicode here, this is for binary data
  27                 });
  28                 my $opt = $self->{opt} // {};
  29                 $dbh->do('PRAGMA synchronous = OFF') if !$opt->{fsync};
  30                 if (my $s = $opt->{cache_size}) {
  31                         $dbh->do("PRAGMA cache_size = $s");
  32                 }
  33                 $dbh->do('PRAGMA journal_mode = '.
  34                                 ($opt->{journal_mode} // 'WAL'));
  35                 $dbh->do(<<'');
  36 CREATE TABLE IF NOT EXISTS kv (
  37         k VARBINARY PRIMARY KEY NOT NULL,
  38         v VARBINARY NOT NULL,
  39         UNIQUE (k)
  40 )
  41
  42                 $dbh;
  43         }
  44 }
  45
  46 sub new {
  47         my ($cls, $dir, $base, $opt) = @_;
  48         my $self = bless { opt => $opt }, $cls;
  49         make_path($dir) if defined($dir) && !-d $dir;
  50         $dir //= $self->{"tmp$$.$self"} = tempdir("skv.$$-XXXX", TMPDIR => 1);
  51         $base //= '';
  52         my $f = $self->{filename} = "$dir/$base.sqlite3";
  53         $self->{lock_path} = $opt->{lock_path} // "$dir/$base.flock";
  54         unless (-f $f) {
  55                 open my $fh, '+>>', $f or die "failed to open $f: $!";
  56                 PublicInbox::Spawn::nodatacow_fd(fileno($fh));
  57         }
  58         $self;
  59 }
  60
  61 sub index_values {
  62         my ($self) = @_;
  63         my $lock = $self->lock_for_scope_fast;
  64         $self->dbh($lock)->do('CREATE INDEX IF NOT EXISTS idx_v ON kv (v)');
  65 }
  66
  67 sub set_maybe {
  68         my ($self, $key, $val, $lock) = @_;
  69         $lock //= $self->lock_for_scope_fast;
  70         my $e = $self->{dbh}->prepare_cached(<<'')->execute($key, $val);
  71 INSERT OR IGNORE INTO kv (k,v) VALUES (?, ?)
  72
  73         $e == 0 ? undef : $e;
  74 }
  75
  76 # caller calls sth->fetchrow_array
  77 sub each_kv_iter {
  78         my ($self) = @_;
  79         my $sth = $self->{dbh}->prepare_cached(<<'', undef, 1);
  80 SELECT k,v FROM kv
  81
  82         $sth->execute;
  83         $sth
  84 }
  85
  86 sub keys {
  87         my ($self, @pfx) = @_;
  88         my $sql = 'SELECT k FROM kv';
  89         if (defined $pfx[0]) {
  90                 $sql .= ' WHERE k LIKE ? ESCAPE ?';
  91                 my $anywhere = !!$pfx[1];
  92                 $pfx[1] = '\\';
  93                 $pfx[0] =~ s/([%_\\])/\\$1/g; # glob chars
  94                 $pfx[0] .= '%';
  95                 substr($pfx[0], 0, 0, '%') if $anywhere;
  96         } else {
  97                 @pfx = (); # [0] may've been undef
  98         }
  99         map { $_->[0] } @{$self->dbh->selectall_arrayref($sql, undef, @pfx)};
 100 }
 101
 102 sub delete_by_val {
 103         my ($self, $val, $lock) = @_;
 104         $lock //= $self->lock_for_scope_fast;
 105         $self->{dbh}->prepare_cached(<<'')->execute($val) + 0;
 106 DELETE FROM kv WHERE v = ?
 107
 108 }
 109
 110 sub replace_values {
 111         my ($self, $oldval, $newval, $lock) = @_;
 112         $lock //= $self->lock_for_scope_fast;
 113         $self->{dbh}->prepare_cached(<<'')->execute($newval, $oldval) + 0;
 114 UPDATE kv SET v = ? WHERE v = ?
 115
 116 }
 117
 118 sub set {
 119         my ($self, $key, $val) = @_;
 120         if (defined $val) {
 121                 my $e = $self->{dbh}->prepare_cached(<<'')->execute($key, $val);
 122 INSERT OR REPLACE INTO kv (k,v) VALUES (?,?)
 123
 124                 $e == 0 ? undef : $e;
 125         } else {
 126                 $self->{dbh}->prepare_cached(<<'')->execute($key);
 127 DELETE FROM kv WHERE k = ?
 128
 129         }
 130 }
 131
 132 sub get {
 133         my ($self, $key) = @_;
 134         my $sth = $self->{dbh}->prepare_cached(<<'', undef, 1);
 135 SELECT v FROM kv WHERE k = ?
 136
 137         $sth->execute($key);
 138         $sth->fetchrow_array;
 139 }
 140
 141 sub xchg {
 142         my ($self, $key, $newval, $lock) = @_;
 143         $lock //= $self->lock_for_scope_fast;
 144         my $oldval = get($self, $key);
 145         if (defined $newval) {
 146                 set($self, $key, $newval);
 147         } else {
 148                 $self->{dbh}->prepare_cached(<<'')->execute($key);
 149 DELETE FROM kv WHERE k = ?
 150
 151         }
 152         $oldval;
 153 }
 154
 155 sub count {
 156         my ($self) = @_;
 157         my $sth = $self->{dbh}->prepare_cached(<<'');
 158 SELECT COUNT(k) FROM kv
 159
 160         $sth->execute;
 161         $sth->fetchrow_array;
 162 }
 163
 164 # faster than ->count due to how SQLite works
 165 sub has_entries {
 166         my ($self) = @_;
 167         my @n = $self->{dbh}->selectrow_array('SELECT k FROM kv LIMIT 1');
 168         scalar(@n) ? 1 : undef;
 169 }
 170
 171 sub dbh_release {
 172         my ($self, $lock) = @_;
 173         my $dbh = delete $self->{dbh} or return;
 174         $lock //= $self->lock_for_scope_fast; # may be needed for WAL
 175         %{$dbh->{CachedKids}} = (); # cleanup prepare_cached
 176         $dbh->disconnect;
 177 }
 178
 179 sub DESTROY {
 180         my ($self) = @_;
 181         dbh_release($self);
 182         my $dir = delete $self->{"tmp$$.$self"} or return;
 183         my $tries = 0;
 184         do {
 185                 $! = 0;
 186                 eval { rmtree($dir) };
 187         } while ($@ && $!{ENOENT} && $tries++ < 5);
 188         warn "error removing $dir: $@" if $@;
 189         warn "Took $tries tries to remove $dir\n" if $tries;
 190 }
 191
 192 1;