How to implement connection pool in spark streaming

时间:2021-12-02 00:28:15

在spark streaming的文档里,有这么一段:

def sendPartition(iter):
    # ConnectionPool is a static, lazily initialized pool of connections
    connection = ConnectionPool.getConnection()
    for record in iter:
        connection.send(record)
    # return to the pool for future reuse
    ConnectionPool.returnConnection(connection)

dstream.foreachRDD(lambda rdd: rdd.foreachPartition(sendPartition))

但是怎么让worker得到一个ConectionPool呢?简单的想法是在使用static变量指向一个ConnectionPool。但这里有一个讲究:怎么保证这个ConnectionPool是worker上的,而不是driver上的?

用pyhton为例:

在ConnectionPool.py里实现一个pool

#/usr/bin/python
#connection_pool.py
import
psycopg2 import settings from DBUtils.PooledDB import PooledDB pool = PooledDB(psycopg2, settings.connection_pool_size, host=settings.db_host, database=settings.database, user=settings.db_user, password=settings.db_password)
def getConnection():
return pool.connection()

假设stream的主代码在main.py里,提交spark

spark-submit --py-files connection_pool.py main.py

这样connection_pool.py将被发送到worker执行,main.py里的 sendPartition 在worker节点上执行的时候就可以获得ConnectionPool.getConnection()调用。

这里的关键是明白哪些代码在driver上跑,哪些在worker上跑。