@andrewor14 what I'm suggest is even simpler than that (for this specific data structure). When the stage finishes, you can just remove the associated entries in this hash map. Once a given stages finishes we will never re-execute the a stage with the same stageId AFIAK. Maybe @markhamstra could confirm this.

The thing you are suggesting is needed (I think) but not for this particular case where I believe something simpler will suffice.