def _RecComputeRebalanceSize(mapping, server_id, dspath, subpath):
"""Recursively compute the size of files that need to be moved."""
total = 0
fulldir = utils.JoinPath(dspath, subpath)
for comp in os.listdir(fulldir):
if comp == constants.REBALANCE_DIRECTORY:
continue
# Real path name.
path = utils.JoinPath(fulldir, comp)
# Get basename and extension.
name, unused_extension = os.path.splitext(comp)
if name in COPY_EXCEPTIONS:
logging.info("Skip %s", comp)
continue
if os.path.isdir(path):
total += _RecComputeRebalanceSize(mapping, server_id, dspath,
utils.JoinPath(subpath, comp))
elif os.path.isfile(path):
key = common.MakeDestinationKey(subpath, name)
where = sutils.MapKeyToServer(mapping, key)
if where != server_id:
logging.info("Need to move %s from %d to %d", path, server_id,
where)
total += os.path.getsize(path)
else:
logging.info("File %s stays here", path)
return total
def Get(self, subject):
"""This will create the connection if needed so should not fail."""
filename, directory = common.ResolveSubjectDestination(subject,
self.path_regexes)
key = common.MakeDestinationKey(directory, filename)
try:
return super(SqliteConnectionCache, self).Get(key)
except KeyError:
dirname = utils.JoinPath(self.root_path, directory)
path = utils.JoinPath(dirname, filename) + SQLITE_EXTENSION
dirname = utils.SmartStr(dirname)
path = utils.SmartStr(path)
# Make sure directory exists.
if not os.path.isdir(dirname):
try:
os.makedirs(dirname)
except OSError:
pass
self._EnsureDatabaseExists(path)
connection = SqliteConnection(path)
super(SqliteConnectionCache, self).Put(key, connection)
return connection
def Get(self, subject):
"""This will create the object if needed so should not fail."""
filename, directory = common.ResolveSubjectDestination(
subject, self.path_regexes)
key = common.MakeDestinationKey(directory, filename)
try:
return super(TDBContextCache, self).Get(key)
except KeyError:
root_path = self.RootPath()
dirname = utils.JoinPath(root_path, directory)
path = utils.JoinPath(dirname, filename) + "." + TDB_EXTENSION
dirname = utils.SmartStr(dirname)
path = utils.SmartStr(path)
# Make sure directory exists.
if not os.path.isdir(dirname):
try:
os.makedirs(dirname)
except OSError:
pass
context = TDBContext(path)
super(TDBContextCache, self).Put(key, context)
return context
def Get(self, subject):
"""This will create the object if needed so should not fail."""
filename, directory = common.ResolveSubjectDestination(subject,
self.path_regexes)
key = common.MakeDestinationKey(directory, filename)
try:
return super(RemoteMappingCache, self).Get(key)
except KeyError:
data_server = self.inquirer.MapKey(key)
super(RemoteMappingCache, self).Put(key, data_server)
return data_server
def _RecCopyFiles(rebalance, server_id, dspath, subpath, pool_cache,
removed_list):
"""Recursively send files for moving to the required data server."""
fulldir = utils.JoinPath(dspath, subpath)
mapping = rebalance.mapping
for comp in os.listdir(fulldir):
if comp == constants.REBALANCE_DIRECTORY:
continue
path = utils.JoinPath(fulldir, comp)
name, unused_extension = os.path.splitext(comp)
if name in COPY_EXCEPTIONS:
continue
if os.path.isdir(path):
result = _RecCopyFiles(rebalance, server_id, dspath,
utils.JoinPath(subpath, comp), pool_cache,
removed_list)
if not result:
return False
continue
if not os.path.isfile(path):
continue
key = common.MakeDestinationKey(subpath, name)
where = sutils.MapKeyToServer(mapping, key)
if where != server_id:
server = mapping.servers[where]
addr = server.address
port = server.port
key = (addr, port)
try:
pool = pool_cache[key]
except KeyError:
pool = urllib3.connectionpool.HTTPConnectionPool(addr,
port=port)
pool_cache[key] = pool
logging.info("Need to move %s from %d to %d", key, server_id,
where)
if not _SendFileToServer(pool, path, subpath, comp, rebalance):
return False
removed_list.append(path)
else:
logging.info("File %s stays here", path)
return True
