pymongo deal with AutoReconnect exception

在 MongoDB 套用 replication 的環境下,如果使用 pymongo 的話,需要注意節點切換的問題,這問題會發生在每次都使用同一個連線的情況下,假如 MongoDB 的 master 節點掛掉或是 master 切換到其他節點的時候,pymongo 接下來的第一個 query 會產生 AutoReconnect 的 exception,然後讓使用者自己決定在這個 exception 裡要做什麼。而 pymongo 這樣的處理方式個人覺得真的很爛!如果以一個單純的 MongoDB 節點轉換到套用 replication 的環境來看,有誰一開始就會在 application 裡檢查這個 exception 的?這樣等於程式中所有的 query 都需要檢查這個 exception,說好的 auto failover 呢?我覺得好歹也給個參數設定可以做 retry 的機制這樣才對吧。為了處理這個問題,google 找了一些方式來處理,還好程式中 query 的部分還不算太分散,這邊可以透過幾種方式解決:

  1. MongoDBProxy https://github.com/arngarden/MongoDBProxy

它的做法是透過包裝原本的 pymongo client 來處理 AutoReconnect 的問題,不過這邊還有另外一個問題是像 find() 這類的會產生一個 cursor,而 AutoReconnect 的 exception 只有在 iterate 這個 cursor 的時候才會發生,這種情況用上面的方式就無解。雖然作者在 library 裡也有提供一個 for cursor 的 class,用 DurableCursor 把拿到的 cursor 包裝起來應該就可以,不過這種方式每個查詢的 query 都要包,也是很麻煩就是,我自己也沒用過。

  1. 自己用 decorator 包裝處理

如果你的 MongoDB 的 query 是交由幾個 class 來專門處理的話,可以直接 decorate 這個 class 裡的 method,會省事很多,graceful_auto_reconnect 是處理 reconnect 的主要部分,另外 for_all_static_methodsfor_all_non_static_methods 則是針對不同類型的 method 做包裝,whitelist 是用來指定你想要被 decorate 的 method 名稱,不指定的話,預設就是全部。

MAX_AUTO_RECONNECT_ATTEMPTS = 5


def graceful_auto_reconnect(mongo_op_func):  
    """Gracefully handle a reconnection event."""
    @functools.wraps(mongo_op_func)
    def wrapper(*args, **kwargs):
        for attempt in range(MAX_AUTO_RECONNECT_ATTEMPTS):
            try:
                return mongo_op_func(*args, **kwargs)
            except pymongo.errors.AutoReconnect:
                wait_t = 0.5 * pow(2, attempt)  # exponential back off
                time.sleep(wait_t)
    return wrapper


def for_all_static_methods(decorator, whitelist=None):  
    if white_list is None:
        white_list = []

    def decorate(cls):
        for name, member in vars(cls).items():
            # Static and class methods: do the dark magic
            if isinstance(member, (classmethod, staticmethod)):
                if white_list and name not in white_list:
                    continue
                inner_func = member.__func__
                method_type = type(member)
                decorated = method_type(decorator(inner_func))
                setattr(cls, name, decorated)
                continue
        return cls
    return decorate


def for_all_non_static_methods(decorator, white_list=None):  
    if white_list is None:
        white_list = []

    def decorate(cls):
        for name, member in vars(cls).items():
            # Static and class methods: do the dark magic
            if isinstance(member, (types.FunctionType, types.BuiltinFunctionType)):
                if white_list and name not in white_list:
                    continue
                setattr(cls, name, decorator(member))
        return cls
    return decorate


# Usage example
@for_all_static_methods(graceful_auto_reconnect)
class HBMongoClient(object):  
    pass

References:
Save the Monkey: Reliably Writing to MongoDB
Gracefully handle a PyMongo AutoReconnect
How to decorate class or static methods
Attaching a decorator to all functions within a class
How to wrap every method of a class?

carlcarl

Read more posts by this author.