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);
  15
  16 sub dbh {
  17         my ($self, $lock) = @_;
  18         $self->{dbh} //= do {
  19                 my $f = $self->{filename};
  20                 $lock //= $self->lock_for_scope;
  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 cache_size = '.($opt->{cache_size} || 80000));
  31                 $dbh->do('PRAGMA journal_mode = '.
  32                                 ($opt->{journal_mode} // 'WAL'));
  33                 $dbh->do(<<'');
  34 CREATE TABLE IF NOT EXISTS kv (
  35         k VARBINARY PRIMARY KEY NOT NULL,
  36         v VARBINARY NOT NULL,
  37         UNIQUE (k)
  38 )
  39
  40                 $dbh;
  41         }
  42 }
  43
  44 sub new {
  45         my ($cls, $dir, $base, $opt) = @_;
  46         my $self = bless { opt => $opt }, $cls;
  47         $dir //= $self->{"tmp$$.$self"} = tempdir("skv.$$-XXXX", TMPDIR => 1);
  48         -d $dir or mkdir($dir) or die "mkdir($dir): $!";
  49         $base //= '';
  50         my $f = $self->{filename} = "$dir/$base.sqlite3";
  51         $self->{lock_path} = $opt->{lock_path} // "$dir/$base.flock";
  52         unless (-f $f) {
  53                 open my $fh, '+>>', $f or die "failed to open $f: $!";
  54                 PublicInbox::Spawn::nodatacow_fd(fileno($fh));
  55         }
  56         $self;
  57 }
  58
  59 sub index_values {
  60         my ($self) = @_;
  61         my $lock = $self->lock_for_scope;
  62         $self->dbh($lock)->do('CREATE INDEX IF NOT EXISTS idx_v ON kv (v)');
  63 }
  64
  65 sub set_maybe {
  66         my ($self, $key, $val, $lock) = @_;
  67         $lock //= $self->lock_for_scope;
  68         my $e = $self->{dbh}->prepare_cached(<<'')->execute($key, $val);
  69 INSERT OR IGNORE INTO kv (k,v) VALUES (?, ?)
  70
  71         $e == 0 ? undef : $e;
  72 }
  73
  74 # caller calls sth->fetchrow_array
  75 sub each_kv_iter {
  76         my ($self) = @_;
  77         my $sth = $self->{dbh}->prepare_cached(<<'', undef, 1);
  78 SELECT k,v FROM kv
  79
  80         $sth->execute;
  81         $sth
  82 }
  83
  84 sub delete_by_val {
  85         my ($self, $val, $lock) = @_;
  86         $lock //= $self->lock_for_scope;
  87         $self->{dbh}->prepare_cached(<<'')->execute($val) + 0;
  88 DELETE FROM kv WHERE v = ?
  89
  90 }
  91
  92 sub replace_values {
  93         my ($self, $oldval, $newval, $lock) = @_;
  94         $lock //= $self->lock_for_scope;
  95         $self->{dbh}->prepare_cached(<<'')->execute($newval, $oldval) + 0;
  96 UPDATE kv SET v = ? WHERE v = ?
  97
  98 }
  99
 100 sub set {
 101         my ($self, $key, $val) = @_;
 102         if (defined $val) {
 103                 my $e = $self->{dbh}->prepare_cached(<<'')->execute($key, $val);
 104 INSERT OR REPLACE INTO kv (k,v) VALUES (?,?)
 105
 106                 $e == 0 ? undef : $e;
 107         } else {
 108                 $self->{dbh}->prepare_cached(<<'')->execute($key);
 109 DELETE FROM kv WHERE k = ?
 110
 111         }
 112 }
 113
 114 sub get {
 115         my ($self, $key) = @_;
 116         my $sth = $self->{dbh}->prepare_cached(<<'', undef, 1);
 117 SELECT v FROM kv WHERE k = ?
 118
 119         $sth->execute($key);
 120         $sth->fetchrow_array;
 121 }
 122
 123 sub xchg {
 124         my ($self, $key, $newval, $lock) = @_;
 125         $lock //= $self->lock_for_scope;
 126         my $oldval = get($self, $key);
 127         if (defined $newval) {
 128                 set($self, $key, $newval);
 129         } else {
 130                 $self->{dbh}->prepare_cached(<<'')->execute($key);
 131 DELETE FROM kv WHERE k = ?
 132
 133         }
 134         $oldval;
 135 }
 136
 137 sub count {
 138         my ($self) = @_;
 139         my $sth = $self->{dbh}->prepare_cached(<<'');
 140 SELECT COUNT(k) FROM kv
 141
 142         $sth->execute;
 143         $sth->fetchrow_array;
 144 }
 145
 146 sub DESTROY {
 147         my ($self) = @_;
 148         delete $self->{dbh};
 149         my $dir = delete $self->{"tmp$$.$self"} or return;
 150         my $tries = 0;
 151         do {
 152                 $! = 0;
 153                 eval { rmtree($dir) };
 154         } while ($@ && $!{ENOENT} && $tries++ < 5);
 155         warn "error removing $dir: $@" if $@;
 156         warn "Took $tries tries to remove $dir\n" if $tries;
 157 }
 158
 159 1;