user sessions: get ability to count memcached sessions

rhodecode/lib/user_sessions.py

@@ -1,181 +1,223
 # -*- coding: utf-8 -*-
 
 # Copyright (C) 2017-2017 RhodeCode GmbH
 #
 # This program is free software: you can redistribute it and/or modify
 # it under the terms of the GNU Affero General Public License, version 3
 # (only), as published by the Free Software Foundation.
 #
 # This program is distributed in the hope that it will be useful,
 # but WITHOUT ANY WARRANTY; without even the implied warranty of
 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 # GNU General Public License for more details.
 #
 # You should have received a copy of the GNU Affero General Public License
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 #
 # This program is dual-licensed. If you wish to learn more about the
 # RhodeCode Enterprise Edition, including its added features, Support services,
 # and proprietary license terms, please see https://rhodecode.com/licenses/
 
 import os
+import re
 import time
 import datetime
 import dateutil
+
 from rhodecode.model.db import DbSession, Session
 
 
 class CleanupCommand(Exception):
     pass
 
 
 class BaseAuthSessions(object):
     SESSION_TYPE = None
     NOT_AVAILABLE = 'NOT AVAILABLE'
 
     def __init__(self, config):
         session_conf = {}
         for k, v in config.items():
             if k.startswith('beaker.session'):
                 session_conf[k] = v
         self.config = session_conf
 
     def get_count(self):
         raise NotImplementedError
 
     def get_expired_count(self, older_than_seconds=None):
         raise NotImplementedError
 
     def clean_sessions(self, older_than_seconds=None):
         raise NotImplementedError
 
     def _seconds_to_date(self, seconds):
         return datetime.datetime.utcnow() - dateutil.relativedelta.relativedelta(
             seconds=seconds)
 
 
 class DbAuthSessions(BaseAuthSessions):
     SESSION_TYPE = 'ext:database'
 
     def get_count(self):
         return DbSession.query().count()
 
     def get_expired_count(self, older_than_seconds=None):
         expiry_date = self._seconds_to_date(older_than_seconds)
         return DbSession.query().filter(DbSession.accessed < expiry_date).count()
 
     def clean_sessions(self, older_than_seconds=None):
         expiry_date = self._seconds_to_date(older_than_seconds)
         to_remove = DbSession.query().filter(DbSession.accessed < expiry_date).count()
         DbSession.query().filter(DbSession.accessed < expiry_date).delete()
         Session().commit()
         return to_remove
 
 
 class FileAuthSessions(BaseAuthSessions):
     SESSION_TYPE = 'file sessions'
 
     def _get_sessions_dir(self):
         data_dir = self.config.get('beaker.session.data_dir')
         return data_dir
 
     def _count_on_filesystem(self, path, older_than=0, callback=None):
         value = dict(percent=0, used=0, total=0, items=0, callbacks=0,
                      path=path, text='')
         items_count = 0
         used = 0
         callbacks = 0
         cur_time = time.time()
         for root, dirs, files in os.walk(path):
             for f in files:
                 final_path = os.path.join(root, f)
                 try:
                     mtime = os.stat(final_path).st_mtime
                     if (cur_time - mtime) > older_than:
                         items_count += 1
                         if callback:
                             callback_res = callback(final_path)
                             callbacks += 1
                         else:
                             used += os.path.getsize(final_path)
                 except OSError:
                     pass
         value.update({
             'percent': 100,
             'used': used,
             'total': used,
             'items': items_count,
             'callbacks': callbacks
         })
         return value
 
     def get_count(self):
         try:
             sessions_dir = self._get_sessions_dir()
             items_count = self._count_on_filesystem(sessions_dir)['items']
         except Exception:
             items_count = self.NOT_AVAILABLE
         return items_count
 
     def get_expired_count(self, older_than_seconds=0):
         try:
             sessions_dir = self._get_sessions_dir()
             items_count = self._count_on_filesystem(
                 sessions_dir, older_than=older_than_seconds)['items']
         except Exception:
             items_count = self.NOT_AVAILABLE
         return items_count
 
     def clean_sessions(self, older_than_seconds=0):
         # find . -mtime +60 -exec rm {} \;
 
         sessions_dir = self._get_sessions_dir()
 
         def remove_item(path):
             os.remove(path)
 
         stats = self._count_on_filesystem(
             sessions_dir, older_than=older_than_seconds,
             callback=remove_item)
         return stats['callbacks']
 
 
+
 class MemcachedAuthSessions(BaseAuthSessions):
     SESSION_TYPE = 'ext:memcached'
+    _key_regex = re.compile(r'ITEM (.*_session) \[(.*); (.*)\]')
+
+    def _get_client(self):
+        import memcache
+        client = memcache.Client([self.config.get('beaker.session.url')])
+        return client
+
+    def _get_telnet_client(self, host, port):
+        import telnetlib
+        client = telnetlib.Telnet(host, port, None)
+        return client
+
+    def _run_telnet_cmd(self, client, cmd):
+        client.write("%s\n" % cmd)
+        return client.read_until('END')
+
+    def key_details(self, client, slab_ids, limit=100):
+        """  Return a list of tuples containing keys and details """
+        cmd = 'stats cachedump %s %s'
+        for slab_id in slab_ids:
+            for key in self._key_regex.finditer(
+                    self._run_telnet_cmd(client, cmd % (slab_id, limit))):
+                yield key
 
     def get_count(self):
-        return self.NOT_AVAILABLE
+        client = self._get_client()
+        count = self.NOT_AVAILABLE
+        try:
+            slabs = []
+            for server, slabs_data in client.get_slabs():
+                slabs.extend(slabs_data.keys())
+
+            host, port = client.servers[0].address
+            telnet_client = self._get_telnet_client(host, port)
+            keys = self.key_details(telnet_client, slabs)
+            count = 0
+            for _k in keys:
+                count += 1
+        except Exception:
+            return count
+
+        return count
 
     def get_expired_count(self, older_than_seconds=None):
         return self.NOT_AVAILABLE
 
     def clean_sessions(self, older_than_seconds=None):
         raise CleanupCommand('Cleanup for this session type not yet available')
 
 
 class MemoryAuthSessions(BaseAuthSessions):
     SESSION_TYPE = 'memory'
 
     def get_count(self):
         return self.NOT_AVAILABLE
 
     def get_expired_count(self, older_than_seconds=None):
         return self.NOT_AVAILABLE
 
     def clean_sessions(self, older_than_seconds=None):
         raise CleanupCommand('Cleanup for this session type not yet available')
 
 
 def get_session_handler(session_type):
     types = {
         'file': FileAuthSessions,
         'ext:memcached': MemcachedAuthSessions,
         'ext:database': DbAuthSessions,
         'memory': MemoryAuthSessions
     }
 
     try:
         return types[session_type]
     except KeyError:
         raise ValueError(
             'This type {} is not supported'.format(session_type))