# 直连RADOS上传文件
import rados
import time
import uuid

cluster = rados.Rados(conffile='/etc/ceph/ceph.conf')
cluster.connect()

ioctx = cluster.open_ioctx('cephfs_data')

# 写入文件到集群中
def write_file(file_name):
    # for i in range(1, 10):
    start = time.time()
    fn = file_name + "." + str(uuid.uuid4())
    print("Start writing " + fn)
    with open(file_name, 'rb') as f:
        data = f.read()
        ioctx.write_full(fn, data)
    end = time.time()
    print(fn + " is written. Time: " + str(end - start) + "s")
    print("Speed: " + str(len(data) / (end - start) / 1024 / 1024) + "MB/s")

# 并发测试：开10个线程，每个线程写一个文件
import threading
threads = []
for i in range(1, 11):
    t = threading.Thread(target=write_file, args=("OS.pdf",))
    threads.append(t)
    t.start()

# 列出池中的所有文件名
# ioctx = cluster.open_ioctx('cephfs_data')
# objects = ioctx.list_objects()
# for obj in objects:
#     print(obj.key)
# ioctx.close()

# 读取文件
# ioctx = cluster.open_ioctx('cephfs_data')
# data = ioctx.read("OS.pdf")
# f_download = open("OS_download.pdf", "wb")
# f_download.write(data)
# f_download.close()
# ioctx.close()

import boto
import boto.s3.connection

access_key = 'V446LFHJPNG4Z0DH4MCR'
secret_key = 'rdlvr8qDooUIRbQKJV9dtxhuTwuDByCXi3TPI1nr'
conn = boto.connect_s3(
        aws_access_key_id = access_key,
        aws_secret_access_key = secret_key,
        host = 'cephtest1', port = 80,
        is_secure=False, calling_format = boto.s3.connection.OrdinaryCallingFormat(),
        )

# 创建Bucket
bucket = conn.create_bucket('ceph-s3-bucket')
for bucket in conn.get_all_buckets():
        print ("{name}".format(name = bucket.name))

# 创建pool
ceph osd pool create rbd 32 32
# 验证
ceph osd pool ls
# pool启用rbd
ceph osd pool application enable rbd rbd
# 初始化
rbd pool init -p rbd

# 创建img映像（可以多创）
rbd create rbd-data-img1 --size 1G --pool rbd --image-format 2 --image-feature layering
# 验证
rbd ls --pool rbd -l

# 需要先安装ceph-common
# 前两条可选，如果你配置了epel源那就不需要了。
yum install epel-release
yum install https://mirrors.aliyun.com/ceph/rpm-octopus/el7/noarch/ceph-release-1-1.el7.noarch.rpm -y
yum install ceph-common

rbd -p rbd map rbd-data-img1

命令
osdmaptool {osdmap_filename} --upmap out.txt [–upmap-pool ] [–upmap-max ] [–upmap-deviation ]
其中

upmap-pool ：指定需要优化均衡的存储池名

upmap-max： 指定一次优化的数据条目，默认100，可根据环境业务情况调整该值，一次调整的条目越多，数据迁移会越多，可能对环境业务造成影响。

max-deviation：最大偏差值，默认为0.01（即1％）。如果OSD利用率与平均值之间的差异小于此值，则将被视为完美。

import os
import time

from subprocess import Popen, TimeoutExpired, PIPE
import datetime

UPMAP_MAX = 100 # 最大优化指令行数
UPMAP_DEVIATION = 1 # 最大偏差值，如果OSD利用率与平均值之间的差异小于此值，将被认为完美而不优化。
POOL = "cephfs_data" # pool名称。后期可以做成轮询所有pool或指定的n个pool
ROUND_SECONDS = 1000 # 脚本定时启动时间

FG_COLORS = {
    "black": "30",
    "red": "31",
    "green": "32",
    "yellow": "33",
    "blue": "34",
    "purple": "35",
    "cyan": "36",
    "white": "37",
    "default": "39",
}

BG_COLORS = {
    "black": "40",
    "red": "41",
    "green": "42",
    "yellow": "43",
    "blue": "44",
    "purple": "45",
    "cyan": "46",
    "white": "47",
    "default": "49",
}

LEVEL_INFO = "INFO"
LEVEL_WARNING = "WARNING"
LEVEL_ERROR = "ERROR"
LEVEL_CRITICAL = "CRITICAL"

level_colors = {
    "INFO": "green",
    "WARNING": "yellow",
    "ERROR": "red",
    "CRITICAL": "purple",
}

def get_osd_df(pool):

    osd_df = os.popen("cephadm shell ceph osd df").read()
    osd_df = osd_df.split("\n")[1:]
    osd_df = [line.split(" ") for line in osd_df]
    osd_df = [[word for word in line if word != ""] for line in osd_df]
    osd_df = osd_df[:-3]
    # print("osd_df:")
    # print(osd_df)

    ave_used = 0
    for osd in osd_df:
        ave_used += float(osd[16])
    ave_used /= len(osd_df)
    log(f"pool {pool} OSD Average used: " + str(ave_used), level=LEVEL_INFO)
    #print("Average used: " + str(ave_used))

    osd_to_upmap = []
    for osd in osd_df:
        if osd[19] != "UP":
            continue
        # print("osd: " + str(osd))
        if abs(float(osd[16]) - ave_used) > UPMAP_DEVIATION:
            osd_to_upmap.append(osd[0])
   # print("OSD to upmap: " + str(osd_to_upmap))

    if len(osd_to_upmap) > 0:
        log(f"Found unbalanced OSD in pool {pool}: " + str(osd_to_upmap), level=LEVEL_INFO)

        # os.popen("cephadm shell ceph osd getmap -o osd.map").read()

        # p = Popen(["cephadm", "shell"], stdin=PIPE, stdout=PIPE, stderr=PIPE)
        # p.stdin.write(b"ceph osd getmap -o osd.map\n")
        # p.stdin.write(b"osdmaptool osd.map --upmap " + b" ".join(osd_to_upmap) + b" --upmap-pool " + bytes(POOL, encoding="utf-8") + b" --upmap-max " + bytes(str(UPMAP_MAX), encoding="utf-8") + b" --upmap-deviation " + bytes(str(UPMAP_DEVIATION), encoding="utf-8") + b"\n")
        # p.stdin.write(b"exit\n")
        # p.stdin.close()
        # print(p.stdout.read().decode("utf-8"))
        # print(p.stderr.read().decode("utf-8"))

        t = os.popen(f"cephadm shell -- bash -c 'ceph osd getmap -o osd.map; osdmaptool osd.map --upmap t.out --upmap-pool {pool} --upmap-max {UPMAP_MAX} --upmap-deviation {UPMAP_DEVIATION}; cat t.out'").read()
        upmap_cmds = []
        l = t.split("\n")
        for line in l:
            if line.startswith("ceph osd pg-upmap-item"):
                # print(line)
                upmap_cmds.append(line)
        log(f"Upmap commands: " + str(upmap_cmds), level=LEVEL_INFO)
        f_name = f"osd_upmap_{time.time()}"
        if len(upmap_cmds) > 0:
            with open(f_name, "w") as f:
                f.write("\n".join(upmap_cmds))
            # 执行
            t = os.popen(f"chmod +x {f_name} && ./{f_name}").read()
            log(f"Upmap done. Result: " + str(t), level=LEVEL_INFO)
        else:
            log(f"No upmap commands generated, skip.", level=LEVEL_INFO)

# 这里可以不用理解。这只是一个打印log的函数。
def log(
        msg: str,
        level: str = "INFO",
        tag: str = "System",
        fg: str = None,
        bg: str = None,
        max_len: int = 10000):
    
    if len(msg) > max_len:
        msg = msg[:max_len] + "..."
    now = datetime.datetime.now().strftime("%m-%d %H:%M:%S")
    pre = f"[{now}] [{level}] [{tag}]: {msg}"
    if level == "INFO":
        if fg is None:
            fg = FG_COLORS["green"]
        if bg is None:
            bg = BG_COLORS["default"]
    elif level == "WARNING":
        if fg is None:
            fg = FG_COLORS["yellow"]
        if bg is None:
            bg = BG_COLORS["default"]
    elif level == "ERROR":
        if fg is None:
            fg = FG_COLORS["red"]
        if bg is None:
            bg = BG_COLORS["default"]
    elif level == "CRITICAL":
        if fg is None:
            fg = FG_COLORS["purple"]
        if bg is None:
            bg = BG_COLORS["default"]
    
    print(f"\033[{fg};{bg}m{pre}\033[0m")

    

if __name__ == "__main__":
    while True:
        log("Balance check start", level=LEVEL_INFO, bg=BG_COLORS["yellow"])
        get_osd_df(pool=POOL)
        log("Balance check end", level=LEVEL_INFO, bg=BG_COLORS["yellow"])
        time.sleep(ROUND_SECONDS)