six locks: Add a callback to six_lock_type()

This allows for functionality like mutex_lock_killable(), but in a
generic manner: to be used for deadlock avoidance in bcachefs.

Also add a function for waking up all waiters on a lock, causing them to
recheck if they should bail out.

Signed-off-by: Kent Overstreet <kent.overstreet@gmail.com>

2 files changed, 55 insertions, 13 deletions

diff --git a/include/linux/six.h b/include/linux/six.h
index 0fb1b2f49345..a16e94f482e9 100644
--- a/include/linux/six.h
+++ b/include/linux/six.h
@@ -115,6 +115,8 @@ struct six_lock {
 #endif
 };
 
+typedef int (*six_lock_should_sleep_fn)(struct six_lock *lock, void *);
+
 static __always_inline void __six_lock_init(struct six_lock *lock,
 					    const char *name,
 					    struct lock_class_key *key)
@@ -141,7 +143,7 @@ do {									\
 #define __SIX_LOCK(type)						\
 bool six_trylock_##type(struct six_lock *);				\
 bool six_relock_##type(struct six_lock *, u32);				\
-void six_lock_##type(struct six_lock *);				\
+int six_lock_##type(struct six_lock *, six_lock_should_sleep_fn, void *);\
 void six_unlock_##type(struct six_lock *);
 
 __SIX_LOCK(read)
@@ -167,14 +169,15 @@ static inline bool six_trylock_type(struct six_lock *lock, enum six_lock_type ty
 }
 
 static inline bool six_relock_type(struct six_lock *lock, enum six_lock_type type,
-		     unsigned seq)
+				   unsigned seq)
 {
 	SIX_LOCK_DISPATCH(type, six_relock, lock, seq);
 }
 
-static inline void six_lock_type(struct six_lock *lock, enum six_lock_type type)
+static inline int six_lock_type(struct six_lock *lock, enum six_lock_type type,
+				six_lock_should_sleep_fn should_sleep_fn, void *p)
 {
-	SIX_LOCK_DISPATCH(type, six_lock, lock);
+	SIX_LOCK_DISPATCH(type, six_lock, lock, should_sleep_fn, p);
 }
 
 static inline void six_unlock_type(struct six_lock *lock, enum six_lock_type type)
@@ -189,4 +192,6 @@ bool six_trylock_convert(struct six_lock *, enum six_lock_type,
 
 void six_lock_increment(struct six_lock *, enum six_lock_type);
 
+void six_lock_wakeup_all(struct six_lock *);
+
 #endif /* _LINUX_SIX_H */
diff --git a/kernel/locking/six.c b/kernel/locking/six.c
index 3d863a9b108d..49d46ed2e18e 100644
--- a/kernel/locking/six.c
+++ b/kernel/locking/six.c
@@ -267,15 +267,21 @@ static inline bool six_optimistic_spin(struct six_lock *lock, enum six_lock_type
 #endif
 
 noinline
-static void __six_lock_type_slowpath(struct six_lock *lock, enum six_lock_type type)
+static int __six_lock_type_slowpath(struct six_lock *lock, enum six_lock_type type,
+				    six_lock_should_sleep_fn should_sleep_fn, void *p)
 {
 	const struct six_lock_vals l[] = LOCK_VALS;
 	union six_lock_state old, new;
 	struct six_lock_waiter wait;
+	int ret = 0;
 	u64 v;
 
+	ret = should_sleep_fn ? should_sleep_fn(lock, p) : 0;
+	if (ret)
+		return ret;
+
 	if (six_optimistic_spin(lock, type))
-		return;
+		return 0;
 
 	lock_contended(&lock->dep_map, _RET_IP_);
 
@@ -292,6 +298,10 @@ static void __six_lock_type_slowpath(struct six_lock *lock, enum six_lock_type t
 			raw_spin_unlock(&lock->wait_lock);
 		}
 
+		ret = should_sleep_fn ? should_sleep_fn(lock, p) : 0;
+		if (ret)
+			break;
+
 		v = READ_ONCE(lock->state.v);
 		do {
 			new.v = old.v = v;
@@ -311,7 +321,8 @@ static void __six_lock_type_slowpath(struct six_lock *lock, enum six_lock_type t
 		schedule();
 	}
 
-	six_set_owner(lock, type, old);
+	if (!ret)
+		six_set_owner(lock, type, old);
 
 	__set_current_state(TASK_RUNNING);
 
@@ -320,18 +331,28 @@ static void __six_lock_type_slowpath(struct six_lock *lock, enum six_lock_type t
 		list_del_init(&wait.list);
 		raw_spin_unlock(&lock->wait_lock);
 	}
+
+	return ret;
 }
 
 __always_inline
-static void __six_lock_type(struct six_lock *lock, enum six_lock_type type)
+static int __six_lock_type(struct six_lock *lock, enum six_lock_type type,
+			   six_lock_should_sleep_fn should_sleep_fn, void *p)
 {
+	int ret;
+
 	if (type != SIX_LOCK_write)
 		six_acquire(&lock->dep_map, 0);
 
-	if (!do_six_trylock_type(lock, type))
-		__six_lock_type_slowpath(lock, type);
+	ret = do_six_trylock_type(lock, type) ? 0
+		: __six_lock_type_slowpath(lock, type, should_sleep_fn, p);
 
-	lock_acquired(&lock->dep_map, _RET_IP_);
+	if (ret && type != SIX_LOCK_write)
+		six_release(&lock->dep_map);
+	if (!ret)
+		lock_acquired(&lock->dep_map, _RET_IP_);
+
+	return ret;
 }
 
 static inline void six_lock_wakeup(struct six_lock *lock,
@@ -417,9 +438,10 @@ bool six_relock_##type(struct six_lock *lock, u32 seq)			\
 }									\
 EXPORT_SYMBOL_GPL(six_relock_##type);					\
 									\
-void six_lock_##type(struct six_lock *lock)				\
+int six_lock_##type(struct six_lock *lock,				\
+		    six_lock_should_sleep_fn should_sleep_fn, void *p)	\
 {									\
-	__six_lock_type(lock, SIX_LOCK_##type);				\
+	return __six_lock_type(lock, SIX_LOCK_##type, should_sleep_fn, p);\
 }									\
 EXPORT_SYMBOL_GPL(six_lock_##type);					\
 									\
@@ -514,3 +536,18 @@ void six_lock_increment(struct six_lock *lock, enum six_lock_type type)
 	}
 }
 EXPORT_SYMBOL_GPL(six_lock_increment);
+
+void six_lock_wakeup_all(struct six_lock *lock)
+{
+	struct six_lock_waiter *w;
+
+	raw_spin_lock(&lock->wait_lock);
+
+	list_for_each_entry(w, &lock->wait_list[0], list)
+		wake_up_process(w->task);
+	list_for_each_entry(w, &lock->wait_list[1], list)
+		wake_up_process(w->task);
+
+	raw_spin_unlock(&lock->wait_lock);
+}
+EXPORT_SYMBOL_GPL(six_lock_wakeup_all);