diff --git a/drivers/gpu/drm/drm_atomic.c b/drivers/gpu/drm/drm_atomic.c
index e4879d3..86e89db 100644
--- a/drivers/gpu/drm/drm_atomic.c
+++ b/drivers/gpu/drm/drm_atomic.c
@@ -144,15 +144,7 @@
 		if (!connector)
 			continue;
 
-		/*
-		 * FIXME: Nonblocking commits can race with connector unplugging and
-		 * there's currently nothing that prevents cleanup up state for
-		 * deleted connectors. As long as the callback doesn't look at
-		 * the connector we'll be fine though, so make sure that's the
-		 * case by setting all connector pointers to NULL.
-		 */
-		state->connector_states[i]->connector = NULL;
-		connector->funcs->atomic_destroy_state(NULL,
+		connector->funcs->atomic_destroy_state(connector,
 						       state->connector_states[i]);
 		state->connectors[i] = NULL;
 		state->connector_states[i] = NULL;
@@ -1168,6 +1160,8 @@
 {
 	struct drm_crtc_state *crtc_state;
 
+	if (crtc)
+		drm_connector_reference(conn_state->connector);
 	if (conn_state->crtc && conn_state->crtc != crtc) {
 		crtc_state = drm_atomic_get_existing_crtc_state(conn_state->state,
 								conn_state->crtc);
@@ -1185,6 +1179,8 @@
 			1 << drm_connector_index(conn_state->connector);
 	}
 
+	if (conn_state->crtc)
+		drm_connector_unreference(conn_state->connector);
 	conn_state->crtc = crtc;
 
 	if (crtc)
diff --git a/drivers/gpu/drm/drm_atomic_helper.c b/drivers/gpu/drm/drm_atomic_helper.c
index 92e11a2..997fd21 100644
--- a/drivers/gpu/drm/drm_atomic_helper.c
+++ b/drivers/gpu/drm/drm_atomic_helper.c
@@ -2762,6 +2762,8 @@
 					    struct drm_connector_state *state)
 {
 	memcpy(state, connector->state, sizeof(*state));
+	if (state->crtc)
+		drm_connector_reference(connector);
 }
 EXPORT_SYMBOL(__drm_atomic_helper_connector_duplicate_state);
 
@@ -2889,6 +2891,8 @@
 	 * state will automatically do the right thing if code is ever added
 	 * to this function.
 	 */
+	if (state->crtc)
+		drm_connector_unreference(state->connector);
 }
 EXPORT_SYMBOL(__drm_atomic_helper_connector_destroy_state);