It's a bit weird to set `context.interrupted` here, since that is meant to be set from outside by Spark only if we actually get interrupted. How about only setting `success = true` and then, above this, changing `while (!context.interrupted)` to `while (!context.interrupted && !success)`?