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 _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
def MapKey(self, key): """Return the data server responsible for a given key.""" sid = sutils.MapKeyToServer(self.mapping, key) return self.servers[sid]
