summaryrefslogtreecommitdiffstats
path: root/lib/pstore.rb
blob: 93795495fca42e19de3ca91aea6ac469b6036346 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
#
# How to use:
#
# db = PStore.new("/tmp/foo")
# db.transaction do
#   p db.roots
#   ary = db["root"] = [1,2,3,4]
#   ary[0] = [1,1.5]
# end

# db.transaction do
#   p db["root"]
# end

require "ftools"
require "digest/md5"

class PStore
  class Error < StandardError
  end

  def initialize(file)
    dir = File::dirname(file)
    unless File::directory? dir
      raise PStore::Error, format("directory %s does not exist", dir)
    end
    if File::exist? file and not File::readable? file
      raise PStore::Error, format("file %s not readable", file)
    end
    @transaction = false
    @filename = file
    @abort = false
  end

  def in_transaction
    raise PStore::Error, "not in transaction" unless @transaction
  end
  private :in_transaction

  def [](name)
    in_transaction
    @table[name]
  end
  def fetch(name, default=PStore::Error)
    unless @table.key? name
      if default==PStore::Error
	raise PStore::Error, format("undefined root name `%s'", name)
      else
	default
      end
    end
    self[name]
  end
  def []=(name, value)
    in_transaction
    @table[name] = value
  end
  def delete(name)
    in_transaction
    @table.delete name
  end

  def roots
    in_transaction
    @table.keys
  end
  def root?(name)
    in_transaction
    @table.key? name
  end
  def path
    @filename
  end

  def commit
    in_transaction
    @abort = false
    throw :pstore_abort_transaction
  end
  def abort
    in_transaction
    @abort = true
    throw :pstore_abort_transaction
  end

  def transaction(read_only=false)
    raise PStore::Error, "nested transaction" if @transaction
    begin
      @transaction = true
      value = nil
      backup = @filename+"~"
      begin
	file = File::open(@filename, read_only ? "rb" : "rb+")
	orig = true
      rescue Errno::ENOENT
	raise if read_only
	file = File::open(@filename, "wb+")
      end
      file.flock(read_only ? File::LOCK_SH : File::LOCK_EX)
      if read_only
	@table = Marshal::load(file)
      elsif orig and (content = file.read) != ""
	@table = Marshal::load(content)
	size = content.size
	md5 = Digest::MD5.digest(content)
	content = nil		# unreference huge data
      else
	@table = {}
      end
      begin
	catch(:pstore_abort_transaction) do
	  value = yield(self)
	end
      rescue Exception
	@abort = true
	raise
      ensure
	if !read_only and !@abort
	  file.rewind
	  content = Marshal::dump(@table)
	  if !md5 || size != content.size || md5 != Digest::MD5.digest(content)
	    File::copy @filename, backup
	    begin
	      file.write(content)
	      file.truncate(file.pos)
	      content = nil		# unreference huge data
	    rescue
	      File::rename backup, @filename if File::exist?(backup)
	      raise
	    end
	  end
	end
	if @abort and !orig
	  File.unlink(@filename)
	end
	@abort = false
      end
    ensure
      @table = nil
      @transaction = false
      file.close if file
    end
    value
  end
end

if __FILE__ == $0
  db = PStore.new("/tmp/foo")
  db.transaction do
    p db.roots
    ary = db["root"] = [1,2,3,4]
    ary[1] = [1,1.5]
  end

  1000.times do
    db.transaction do
      db["root"][0] += 1
      p db["root"][0]
    end
  end

  db.transaction(true) do
    p db["root"]
  end
end