4 files changed, 52 insertions, 0 deletions

diff --git a/lib/dbwrap/dbwrap.c b/lib/dbwrap/dbwrap.c
index d46044f0c4..835bd599ef 100644
--- a/lib/dbwrap/dbwrap.c
+++ b/lib/dbwrap/dbwrap.c
@@ -231,6 +231,33 @@ struct db_record *dbwrap_try_fetch_locked(struct db_context *db,
 		? db->try_fetch_locked : db->fetch_locked);
 }
 
+struct db_record *dbwrap_fetch_locked_timeout(struct db_context *db,
+					      TALLOC_CTX *mem_ctx,
+					      TDB_DATA key,
+					      unsigned int timeout)
+{
+	struct db_record *rec;
+	struct dbwrap_lock_order_state *lock_order;
+	TALLOC_CTX *frame = talloc_stackframe();
+
+	lock_order = dbwrap_check_lock_order(db, frame);
+	if (lock_order == NULL) {
+		TALLOC_FREE(frame);
+		return NULL;
+	}
+	rec = db->fetch_locked_timeout
+		? db->fetch_locked_timeout(db, mem_ctx, key, timeout)
+		: db->fetch_locked(db, mem_ctx, key);
+	if (rec == NULL) {
+		TALLOC_FREE(frame);
+		return NULL;
+	}
+	(void)talloc_steal(rec, lock_order);
+	rec->db = db;
+	TALLOC_FREE(frame);
+	return rec;
+}
+
 struct db_context *dbwrap_record_get_db(struct db_record *rec)
 {
 	return rec->db;
diff --git a/lib/dbwrap/dbwrap.h b/lib/dbwrap/dbwrap.h
index 23a43da019..0b19396b98 100644
--- a/lib/dbwrap/dbwrap.h
+++ b/lib/dbwrap/dbwrap.h
@@ -43,6 +43,11 @@ struct db_record *dbwrap_fetch_locked(struct db_context *db,
 struct db_record *dbwrap_try_fetch_locked(struct db_context *db,
 					  TALLOC_CTX *mem_ctx,
 					  TDB_DATA key);
+struct db_record *dbwrap_fetch_locked_timeout(struct db_context *db,
+					      TALLOC_CTX *mem_ctx,
+					      TDB_DATA key,
+					      unsigned int timeout);
+
 struct db_context *dbwrap_record_get_db(struct db_record *rec);
 void dbwrap_set_stored_callback(
 	struct db_context *db,
diff --git a/lib/dbwrap/dbwrap_private.h b/lib/dbwrap/dbwrap_private.h
index af8374dbc9..dfd365dab2 100644
--- a/lib/dbwrap/dbwrap_private.h
+++ b/lib/dbwrap/dbwrap_private.h
@@ -38,6 +38,10 @@ struct db_context {
 	struct db_record *(*try_fetch_locked)(struct db_context *db,
 					      TALLOC_CTX *mem_ctx,
 					      TDB_DATA key);
+	struct db_record *(*fetch_locked_timeout)(struct db_context *db,
+						  TALLOC_CTX *mem_ctx,
+						  TDB_DATA key,
+						  unsigned int timeout);
 	int (*traverse)(struct db_context *db,
 			int (*f)(struct db_record *rec,
 				 void *private_data),
diff --git a/lib/dbwrap/dbwrap_tdb.c b/lib/dbwrap/dbwrap_tdb.c
index fb6841bcb6..99ea62a3a1 100644
--- a/lib/dbwrap/dbwrap_tdb.c
+++ b/lib/dbwrap/dbwrap_tdb.c
@@ -158,6 +158,21 @@ static struct db_record *db_tdb_fetch_locked(
 	return db_tdb_fetch_locked_internal(db, mem_ctx, key);
 }
 
+static struct db_record *db_tdb_fetch_locked_timeout(
+	struct db_context *db, TALLOC_CTX *mem_ctx, TDB_DATA key,
+	unsigned int timeout)
+{
+	struct db_tdb_ctx *ctx = talloc_get_type_abort(db->private_data,
+						       struct db_tdb_ctx);
+
+	db_tdb_log_key("Locking with timeout ", key);
+	if (tdb_chainlock_with_timeout(ctx->wtdb->tdb, key, timeout) != 0) {
+		DEBUG(3, ("tdb_chainlock_with_timeout failed\n"));
+		return NULL;
+	}
+	return db_tdb_fetch_locked_internal(db, mem_ctx, key);
+}
+
 static struct db_record *db_tdb_try_fetch_locked(
 	struct db_context *db, TALLOC_CTX *mem_ctx, TDB_DATA key)
 {
@@ -429,6 +444,7 @@ struct db_context *db_open_tdb(TALLOC_CTX *mem_ctx,
 	db_tdb->id.ino = st.st_ino;
 
 	result->fetch_locked = db_tdb_fetch_locked;
+	result->fetch_locked_timeout = db_tdb_fetch_locked_timeout;
 	result->try_fetch_locked = db_tdb_try_fetch_locked;
 	result->traverse = db_tdb_traverse;
 	result->traverse_read = db_tdb_traverse_read;