Make the flow context based.

Previously we did not have too many go routines, and could get away without a context based approach. But now with addition of taskmanager and some other concurrent features, it would be better if we refactored the code to be context based.