Minor nit: Might be better avoid null if possible, e.g.:

```scala
mapStatuses.get(shuffleId) match {
  case Some(status) => ...
  case _ => ...
}
```