如何在Spark Executor中关闭内核之间的共享单例连接

Question

I am using shared connection between all cores of single executor of Spark. Basically I have created singleton connection object in order to share between cores of single executor so that it will be shared between cores and there will be only 1 connection per executor.

object SingletonConnection {

private var connection: Connection = null

def getConnection(url: String, username: String, password: String): Connection = synchronized {
if (connection == null) {
  connection = DriverManager.getConnection(url, username, password)
}
connection
}
}

Spark executor code:

dataFrame.foreachPartition { batch =>
  if (batch.nonEmpty) {
    lazy val dbConnection = SingletonConnection
    val dbc = dbConnection.getConnection(url, user, password)

    // do some operatoins


          st.addBatch()
        }
        st.executeBatch()
      }
    }
    catch {
      case exec: BatchUpdateException =>
        var ex: SQLException = exec
        while (ex != null) {
          ex.printStackTrace()
          ex = ex.getNextException
        }
        throw exec
    }

  }
}

Problem here is , I cannot close the connection. Since I will not know when particular core finishes its execution. If i close connection in finally, as soon as one core finishes its task it closes the connection and that causes all other cores to stop since shared connection is closed.

Since I am not closing the connection here, the connection remains open even after the task is finished. How can I make this process work so that I should be able to close the connection ONLY AFTER ALL CORES HAVE FINISHED THEIR TASKS.

Answer 1

I implemented it using Java, so I can just give you some clue.

In SingletonConnection class I created a thread-safe accumulator. Each time the connection is opened, the accumulator is incremented by one. And each time befor closing the connection, the accumulator is decremented by one and check if the accumulator is equals to zero. When the accumulator equals to zero, then you can close the connection.

This won't close connection when other runnning threads are still using the connection. But this will let you create more connections than you thought(the number of partitions).