Repository: iterative/dvc
Pull request #1128 (2018-09-18 00:43:53)

state: limit number of entries


Patch:
diff --git a/dvc/config.py b/dvc/config.py
index 15091ebc9..d8abad665 100644
--- a/dvc/config.py
+++ b/dvc/config.py
@@ -45,6 +45,14 @@ def to_bool(val):
     return val.lower() == 'true'
 
 
+def is_whole(val):
+    return int(val) >= 0
+
+
+def is_percent(val):
+    return int(val) >= 0 and int(val) <= 100
+
+
 class Config(object):
     CONFIG = 'config'
     CONFIG_LOCAL = 'config.local'
@@ -149,10 +157,19 @@ class Config(object):
         Optional(SECTION_REMOTE_TIMEOUT): Use(int),
     }
 
+    SECTION_STATE = 'state'
+    SECTION_STATE_ROW_LIMIT = 'row_limit'
+    SECTION_STATE_ROW_CLEANUP_QUOTA = 'row_cleanup_quota'
+    SECTION_STATE_SCHEMA = {
+        Optional(SECTION_STATE_ROW_LIMIT): And(Use(int), is_whole),
+        Optional(SECTION_STATE_ROW_CLEANUP_QUOTA): And(Use(int), is_percent),
+    }
+
     SCHEMA = {
         Optional(SECTION_CORE, default={}): SECTION_CORE_SCHEMA,
         Optional(Regex(SECTION_REMOTE_REGEX)): SECTION_REMOTE_SCHEMA,
         Optional(SECTION_CACHE, default={}): SECTION_CACHE_SCHEMA,
+        Optional(SECTION_STATE, default={}): SECTION_STATE_SCHEMA,
 
         # backward compatibility
         Optional(SECTION_AWS, default={}): SECTION_AWS_SCHEMA,
diff --git a/dvc/project.py b/dvc/project.py
index 141031993..9c60b4a4e 100644
--- a/dvc/project.py
+++ b/dvc/project.py
@@ -38,7 +38,7 @@ class Project(object):
         self.lock = Lock(self.dvc_dir)
         # NOTE: storing state and link_state in the repository itself to avoid
         # any possible state corruption in 'shared cache dir' scenario.
-        self.state = State(self)
+        self.state = State(self, self.config._config)
 
         core = self.config._config[Config.SECTION_CORE]
         self.logger = Logger(core.get(Config.SECTION_CORE_LOGLEVEL, None))
diff --git a/dvc/state.py b/dvc/state.py
index 5eda5330f..337645c3b 100644
--- a/dvc/state.py
+++ b/dvc/state.py
@@ -3,6 +3,7 @@ import time
 import sqlite3
 import nanotime
 
+from dvc.config import Config
 from dvc.system import System
 from dvc.utils import file_md5, remove
 from dvc.exceptions import DvcException
@@ -21,16 +22,32 @@ class State(object):
                          "md5 TEXT NOT NULL, " \
                          "timestamp TEXT NOT NULL"
 
+    STATE_INFO_TABLE = 'state_info'
+    STATE_INFO_TABLE_LAYOUT = 'count INTEGER'
+    STATE_INFO_ROW = 1
+
     LINK_STATE_TABLE = 'link_state'
     LINK_STATE_TABLE_LAYOUT = "path TEXT PRIMARY KEY, " \
                               "inode INTEGER NOT NULL, " \
                               "mtime TEXT NOT NULL"
 
-    def __init__(self, project):
+    STATE_ROW_LIMIT = 10000000
+    STATE_ROW_CLEANUP_QUOTA = 50
+
+    def __init__(self, project, config):
         self.project = project
         self.dvc_dir = project.dvc_dir
         self.root_dir = project.root_dir
 
+        self.row_limit = 100
+        self.row_cleanup_quota = 50
+
+        c = config.get(Config.SECTION_STATE, {})
+        self.row_limit = c.get(Config.SECTION_STATE_ROW_LIMIT,
+                               self.STATE_ROW_LIMIT)
+        self.row_cleanup_quota = c.get(Config.SECTION_STATE_ROW_CLEANUP_QUOTA,
+                                       self.STATE_ROW_CLEANUP_QUOTA)
+
         if not self.dvc_dir:
             self.state_file = None
             return
@@ -38,6 +55,7 @@ class State(object):
         self.state_file = os.path.join(self.dvc_dir, self.STATE_FILE)
         self.db = None
         self.c = None
+        self.inserts = 0
 
     def __enter__(self):
         self.load()
@@ -67,6 +85,8 @@ class State(object):
         retries = 1
         while True:
             assert self.db is None
+            assert self.c is None
+            assert self.inserts == 0
             self.db = sqlite3.connect(self.state_file)
             self.c = self.db.cursor()
 
@@ -77,14 +97,23 @@ class State(object):
                 cmd = "CREATE TABLE IF NOT EXISTS {} ({})"
                 self.c.execute(cmd.format(self.STATE_TABLE,
                                           self.STATE_TABLE_LAYOUT))
+                self.c.execute(cmd.format(self.STATE_INFO_TABLE,
+                                          self.STATE_INFO_TABLE_LAYOUT))
                 self.c.execute(cmd.format(self.LINK_STATE_TABLE,
                                           self.LINK_STATE_TABLE_LAYOUT))
+
+                cmd = "INSERT OR IGNORE INTO {} (count) SELECT 0 " \
+                      "WHERE NOT EXISTS (SELECT * FROM {})"
+                self.c.execute(cmd.format(self.STATE_INFO_TABLE,
+                                          self.STATE_INFO_TABLE))
+
                 return
             except sqlite3.DatabaseError:
-                self.db.close()
                 self.c.close()
+                self.db.close()
                 self.db = None
                 self.c = None
+                self.inserts = 0
                 if retries > 0:
                     os.unlink(self.state_file)
                     retries -= 1
@@ -93,11 +122,48 @@ class State(object):
 
     def dump(self):
         assert self.db is not None
+
+        cmd = "SELECT count from {} WHERE rowid={}"
+        self.c.execute(cmd.format(self.STATE_INFO_TABLE,
+                                  self.STATE_INFO_ROW))
+        ret = self.c.fetchall()
+        assert len(ret) == 1
+        assert len(ret[0]) == 1
+        count = ret[0][0] + self.inserts
+
+        if count > self.row_limit:
+            msg = "Cleaning up state. This might take a while."
+            self.project.logger.warn(msg)
+
+            delete = (count - self.row_limit)
+            delete += int(self.row_limit * (self.row_cleanup_quota/100.))
+            cmd = "DELETE FROM {} WHERE timestamp IN (" \
+                  "SELECT timestamp FROM {} ORDER BY timestamp ASC LIMIT {});"
+            self.c.execute(cmd.format(self.STATE_TABLE,
+                                      self.STATE_TABLE,
+                                      delete))
+
+            self.c.execute("VACUUM")
+
+            cmd = "SELECT COUNT(*) FROM {}"
+
+            self.c.execute(cmd.format(self.STATE_TABLE))
+            ret = self.c.fetchall()
+            assert len(ret) == 1
+            assert len(ret[0]) == 1
+            count = ret[0][0]
+
+        cmd = "UPDATE {} SET count = {} WHERE rowid = {}"
+        self.c.execute(cmd.format(self.STATE_INFO_TABLE,
+                                  count,
+                                  self.STATE_INFO_ROW))
+
         self.db.commit()
         self.c.close()
         self.db.close()
         self.db = None
         self.c = None
+        self.inserts = 0
 
     @staticmethod
     def mtime(path):
@@ -137,6 +203,7 @@ class State(object):
                                       mtime,
                                       md5,
                                       int(nanotime.timestamp(time.time()))))
+            self.inserts += 1
         else:
             assert len(ret) == 1
             assert len(ret[0]) == 4


Test patch:
diff --git a/tests/test_state.py b/tests/test_state.py
index ce9bcd5be..fa7894db2 100644
--- a/tests/test_state.py
+++ b/tests/test_state.py
@@ -16,7 +16,7 @@ class TestState(TestDvc):
         mtime = os.path.getmtime(path)
         inode = System.inode(path)
 
-        state = State(self.dvc)
+        state = State(self.dvc, self.dvc.config._config)
 
         with state:
             entry_md5 = state.update(path)