Fix potential deadlock in the table manager (#5472)

* Fix potential deadlock in the table manager

* Ensure second lock is released on return

* use t.TempDir instead of os.TempDir

* fix broken test

Co-authored-by: Sandeep Sukhani <sandeep.d.sukhani@gmail.com>

pkg/storage/stores/shipper/downloads/table_manager.go

@@ -168,6 +168,8 @@ func (tm *TableManager) getOrCreateTable(tableName string) (Table, error) {
 
 	if !ok {
 		tm.tablesMtx.Lock()
+		defer tm.tablesMtx.Unlock()
+
 		// check if some other competing goroutine got the lock before us and created the table, use it if so.
 		table, ok = tm.tables[tableName]
 		if !ok {
@@ -183,7 +185,6 @@ func (tm *TableManager) getOrCreateTable(tableName string) (Table, error) {
 			table = NewTable(tableName, filepath.Join(tm.cfg.CacheDir, tableName), tm.indexStorageClient, tm.boltIndexClient, tm.metrics)
 			tm.tables[tableName] = table
 		}
-		tm.tablesMtx.Unlock()
 	}
 
 	return table, nil

pkg/storage/stores/shipper/downloads/table_manager_test.go

@@ -3,6 +3,7 @@ package downloads
 import (
 	"context"
 	"fmt"
+	"os"
 	"path/filepath"
 	"testing"
 	"time"
@@ -35,29 +36,51 @@ func buildTestTableManager(t *testing.T, path string) (*TableManager, stopFunc)
 }
 
 func TestTableManager_QueryPages(t *testing.T) {
-	tempDir := t.TempDir()
-
-	objectStoragePath := filepath.Join(tempDir, objectsStorageDirName)
+	t.Run("QueryPages", func(t *testing.T) {
+		tempDir := t.TempDir()
+		objectStoragePath := filepath.Join(tempDir, objectsStorageDirName)
 
-	var queries []chunk.IndexQuery
-	for i, name := range []string{"table1", "table2"} {
-		testutil.SetupTable(t, filepath.Join(objectStoragePath, name), testutil.DBsConfig{
-			NumUnCompactedDBs: 5,
-			DBRecordsStart:    i * 1000,
-		}, testutil.PerUserDBsConfig{
-			DBsConfig: testutil.DBsConfig{
+		var queries []chunk.IndexQuery
+		for i, name := range []string{"table1", "table2"} {
+			testutil.SetupTable(t, filepath.Join(objectStoragePath, name), testutil.DBsConfig{
 				NumUnCompactedDBs: 5,
-				DBRecordsStart:    i*1000 + 500,
-			},
-			NumUsers: 1,
-		})
-		queries = append(queries, chunk.IndexQuery{TableName: name})
-	}
+				DBRecordsStart:    i * 1000,
+			}, testutil.PerUserDBsConfig{
+				DBsConfig: testutil.DBsConfig{
+					NumUnCompactedDBs: 5,
+					DBRecordsStart:    i*1000 + 500,
+				},
+				NumUsers: 1,
+			})
+			queries = append(queries, chunk.IndexQuery{TableName: name})
+		}
 
-	tableManager, stopFunc := buildTestTableManager(t, tempDir)
-	defer stopFunc()
+		tableManager, stopFunc := buildTestTableManager(t, tempDir)
+		defer stopFunc()
+
+		testutil.TestMultiTableQuery(t, testutil.BuildUserID(0), queries, tableManager, 0, 2000)
+	})
+
+	t.Run("it doesn't deadlock when table create fails", func(t *testing.T) {
+		tempDir := t.TempDir()
+		require.NoError(t, os.Mkdir(filepath.Join(tempDir, "cache"), 0777))
+
+		// This file forces chunk_util.EnsureDirectory to fail. Any write error would cause this
+		// deadlock
+		f, err := os.CreateTemp(filepath.Join(tempDir, "cache"), "not-a-directory")
+		require.NoError(t, err)
+		badTable := filepath.Base(f.Name())
+
+		tableManager, stopFunc := buildTestTableManager(t, tempDir)
+		defer stopFunc()
+
+		err = tableManager.query(context.Background(), badTable, nil, nil)
+		require.Error(t, err)
 
-	testutil.TestMultiTableQuery(t, testutil.BuildUserID(0), queries, tableManager, 0, 2000)
+		// This one deadlocks without the fix
+		err = tableManager.query(context.Background(), badTable, nil, nil)
+		require.Error(t, err)
+	})
 }
 
 func TestTableManager_cleanupCache(t *testing.T) {