protected synchronized int updateWorkerDrainMap() {
        long startTime = System.nanoTime();
        int numRemovedWorkerIds = 0;

        if (drainOpStatusMap.size() > 0) {
            val currentMembership = membershipManager.getCurrentMembership()
                    .stream().map(WorkerInfo::getWorkerId).collect(Collectors.toSet());
            val removeWorkerIds = new ArrayList<String>();

            for (String workerId : drainOpStatusMap.keySet()) {
                if (!currentMembership.contains(workerId)) {
                    removeWorkerIds.add(workerId);
                }
            }
            for (String workerId : removeWorkerIds) {
                drainOpStatusMap.remove(workerId);
            }
            numRemovedWorkerIds = removeWorkerIds.size();
        }

        if (numRemovedWorkerIds > 0) {
            log.info("cleanupWorkerDrainMap removed {} stale workerIds in {} sec",
                    numRemovedWorkerIds, (System.nanoTime() - startTime) / Math.pow(10, 9));
        }

        return numRemovedWorkerIds;
    }
@Test
    public void testUpdateWorkerDrainMap() throws Exception {
        final int numWorkersInDrainMap = 5;
        String workerId;

        // Set up drain status for some of those workers.
        SchedulerManager.DrainOpStatus drainOp;
        for (int ix = 0; ix < numWorkersInDrainMap; ix++) {
            workerId = "worker-" + ix;
            drainOp = SchedulerManager.DrainOpStatus.DrainCompleted;
            if (ix % 2 == 0) {
                drainOp = SchedulerManager.DrainOpStatus.DrainInProgress;
            }
            callGetDrainStatus(workerId, DrainOps.SetDrainStatus, drainOp);
        }
        val oldDrainMap = schedulerManager.getDrainOpsStatusMap();

        final int numWorkersInCurrentMembership = 3;
        List<WorkerInfo> workerInfoList = new LinkedList<>();
        final String workerHostName = "workerHostName";
        final int workerPort = 5000;
        for (int ix = 0; ix < numWorkersInCurrentMembership; ix++) {
            workerId = "worker-" + ix;
            workerInfoList.add(WorkerInfo.of(workerId, workerHostName, workerPort));
        }
        doReturn(workerInfoList).when(membershipManager).getCurrentMembership();

        val numWorkersCleanedUp = schedulerManager.updateWorkerDrainMap();
        Assert.assertEquals(numWorkersCleanedUp, numWorkersInDrainMap - numWorkersInCurrentMembership);

        val newDrainMap = schedulerManager.getDrainOpsStatusMap();
        for (val worker : newDrainMap.keySet()) {
            Assert.assertTrue(oldDrainMap.get(worker) != null);
            WorkerInfo matchedWorker = workerInfoList.stream()
                    .filter(winfo -> worker.equals(winfo.getWorkerId()))
                    .findAny()
                    .orElse(null);
            Assert.assertTrue(matchedWorker != null);
        }
    }