lib/PublicInbox/SharedKV.pm

   1 # Copyright (C) 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                 $dbh->do('PRAGMA journal_mode = '.
  31                                 ($opt->{journal_mode} // 'WAL'));
  32                 $dbh->do(<<'');
  33 CREATE TABLE IF NOT EXISTS kv (
  34         k VARBINARY PRIMARY KEY NOT NULL,
  35         v VARBINARY NOT NULL,
  36         UNIQUE (k)
  37 )
  38
  39                 $self->{dbh} = $dbh;
  40         }
  41 }
  42
  43 sub new {
  44         my ($cls, $dir, $base, $opt) = @_;
  45         my $self = bless { opt => $opt }, $cls;
  46         make_path($dir) if defined($dir) && !-d $dir;
  47         $dir //= $self->{"tmp$$.$self"} = tempdir("skv.$$-XXXX", TMPDIR => 1);
  48         $base //= '';
  49         my $f = $self->{filename} = "$dir/$base.sqlite3";
  50         $self->{lock_path} = $opt->{lock_path} // "$dir/$base.flock";
  51         unless (-s $f) {
  52                 require PublicInbox::Syscall;
  53                 PublicInbox::Syscall::nodatacow_dir($dir); # for journal/shm/wal
  54                 open my $fh, '+>>', $f or die "failed to open $f: $!";
  55         }
  56         $self;
  57 }
  58
  59 sub set_maybe {
  60         my ($self, $key, $val, $lock) = @_;
  61         $lock //= $self->lock_for_scope_fast;
  62         my $e = $self->{dbh}->prepare_cached(<<'')->execute($key, $val);
  63 INSERT OR IGNORE INTO kv (k,v) VALUES (?, ?)
  64
  65         $e == 0 ? undef : $e;
  66 }
  67
  68 # caller calls sth->fetchrow_array
  69 sub each_kv_iter {
  70         my ($self) = @_;
  71         my $sth = $self->{dbh}->prepare_cached(<<'', undef, 1);
  72 SELECT k,v FROM kv
  73
  74         $sth->execute;
  75         $sth
  76 }
  77
  78 sub keys {
  79         my ($self, @pfx) = @_;
  80         my $sql = 'SELECT k FROM kv';
  81         if (defined $pfx[0]) {
  82                 $sql .= ' WHERE k LIKE ? ESCAPE ?';
  83                 my $anywhere = !!$pfx[1];
  84                 $pfx[1] = '\\';
  85                 $pfx[0] =~ s/([%_\\])/\\$1/g; # glob chars
  86                 $pfx[0] .= '%';
  87                 substr($pfx[0], 0, 0, '%') if $anywhere;
  88         } else {
  89                 @pfx = (); # [0] may've been undef
  90         }
  91         map { $_->[0] } @{$self->dbh->selectall_arrayref($sql, undef, @pfx)};
  92 }
  93
  94 sub set {
  95         my ($self, $key, $val) = @_;
  96         if (defined $val) {
  97                 my $e = $self->{dbh}->prepare_cached(<<'')->execute($key, $val);
  98 INSERT OR REPLACE INTO kv (k,v) VALUES (?,?)
  99
 100                 $e == 0 ? undef : $e;
 101         } else {
 102                 $self->{dbh}->prepare_cached(<<'')->execute($key);
 103 DELETE FROM kv WHERE k = ?
 104
 105         }
 106 }
 107
 108 sub get {
 109         my ($self, $key) = @_;
 110         my $sth = $self->{dbh}->prepare_cached(<<'', undef, 1);
 111 SELECT v FROM kv WHERE k = ?
 112
 113         $sth->execute($key);
 114         $sth->fetchrow_array;
 115 }
 116
 117 sub xchg {
 118         my ($self, $key, $newval, $lock) = @_;
 119         $lock //= $self->lock_for_scope_fast;
 120         my $oldval = get($self, $key);
 121         if (defined $newval) {
 122                 set($self, $key, $newval);
 123         } else {
 124                 $self->{dbh}->prepare_cached(<<'')->execute($key);
 125 DELETE FROM kv WHERE k = ?
 126
 127         }
 128         $oldval;
 129 }
 130
 131 sub count {
 132         my ($self) = @_;
 133         my $sth = $self->{dbh}->prepare_cached(<<'');
 134 SELECT COUNT(k) FROM kv
 135
 136         $sth->execute;
 137         $sth->fetchrow_array;
 138 }
 139
 140 # faster than ->count due to how SQLite works
 141 sub has_entries {
 142         my ($self) = @_;
 143         my @n = $self->{dbh}->selectrow_array('SELECT k FROM kv LIMIT 1');
 144         scalar(@n) ? 1 : undef;
 145 }
 146
 147 sub dbh_release {
 148         my ($self, $lock) = @_;
 149         my $dbh = delete $self->{dbh} or return;
 150         $lock //= $self->lock_for_scope_fast; # may be needed for WAL
 151         %{$dbh->{CachedKids}} = (); # cleanup prepare_cached
 152         $dbh->disconnect;
 153 }
 154
 155 sub DESTROY {
 156         my ($self) = @_;
 157         dbh_release($self);
 158         my $dir = delete $self->{"tmp$$.$self"} or return;
 159         my $tries = 0;
 160         do {
 161                 $! = 0;
 162                 eval { rmtree($dir) };
 163         } while ($@ && $!{ENOENT} && $tries++ < 5);
 164         warn "error removing $dir: $@" if $@;
 165         warn "Took $tries tries to remove $dir\n" if $tries;
 166 }
 167
 168 1;