Fix potential race condition with graceful shutdown