Hook cleanup to update_code, because restart is never triggered