Repository: pingcap/tidb
File: pkg/objstore/objectio/writer_test.go (Go)

Code before:
	require.NoError(t, w.Close())
	compressedData := buf.Bytes()

	// default cfg(decode asynchronously)
	synctest.Test(t, func(t *testing.T) {
		pipeReader, pipeWriter := io.Pipe()

		prevRoutineCnt := runtime.NumGoroutine()
		r, err := compressedio.NewReader(compressedio.Zstd, compressedio.DecompressConfig{}, pipeReader)
		synctest.Wait()

		currRoutineCnt := runtime.NumGoroutine()
		require.NoError(t, err)
		require.Greater(t, currRoutineCnt, prevRoutineCnt)

		_, err = pipeWriter.Write(compressedData)
		require.NoError(t, err)
		err = pipeWriter.Close()
		require.NoError(t, err)

		allData, err := io.ReadAll(r)
		require.NoError(t, err)
		require.Equal(t, "data", string(allData))
	})

	// sync decode
	synctest.Test(t, func(t *testing.T) {
		pipeReader, pipeWriter := io.Pipe()

		prevRoutineCnt := runtime.NumGoroutine()
		config := compressedio.DecompressConfig{ZStdDecodeConcurrency: 1}
		r, err := compressedio.NewReader(compressedio.Zstd, config, pipeReader)
		synctest.Wait()

		require.NoError(t, err)
		currRoutineCnt := runtime.NumGoroutine()
		require.Equal(t, prevRoutineCnt, currRoutineCnt)

		// because this is "sync" decode, the Write will block and cause a deadlock. We
		// must use another goroutine. This is another sign that the decode is sync and
		// previous synctest is async.
		go func() {
			_, err := pipeWriter.Write(compressedData)
			require.NoError(t, err)
			err = pipeWriter.Close()
			require.NoError(t, err)
		}()

		allData, err := io.ReadAll(r)
		require.NoError(t, err)
		require.Equal(t, "data", string(allData))

Review comment (bug):
If the test fails before pipeWriter.Close() is called (e.g., at lines 213, 214, or 217), the pipe writer will not be properly closed, which could lead to resource leaks or hanging goroutines in the zstd reader. Consider using defer pipeWriter.Close() after creating the pipe, or wrapping the test logic in a cleanup function to ensure proper cleanup even when assertions fail.

Code after:
	require.NoError(t, err)
	require.NoError(t, w.Close())
	compressedData := buf.Bytes()

	// default cfg(decode asynchronously)
	synctest.Test(t, func(t *testing.T) {
		pipeReader, pipeWriter := io.Pipe()
		defer pipeReader.Close()

		prevRoutineCnt := runtime.NumGoroutine()
		r, err := compressedio.NewReader(compressedio.Zstd, compressedio.DecompressConfig{}, pipeReader)
		synctest.Wait()

		currRoutineCnt := runtime.NumGoroutine()
		require.NoError(t, err)
		require.Greater(t, currRoutineCnt, prevRoutineCnt)

		_, err = pipeWriter.Write(compressedData)
		require.NoError(t, err)
		err = pipeWriter.Close()
		require.NoError(t, err)

		allData, err := io.ReadAll(r)
		require.NoError(t, err)
		require.Equal(t, "data", string(allData))
	})

	// sync decode
	synctest.Test(t, func(t *testing.T) {
		pipeReader, pipeWriter := io.Pipe()
		defer pipeReader.Close()

		prevRoutineCnt := runtime.NumGoroutine()
		config := compressedio.DecompressConfig{ZStdDecodeConcurrency: 1}
		r, err := compressedio.NewReader(compressedio.Zstd, config, pipeReader)
		synctest.Wait()

		require.NoError(t, err)
		currRoutineCnt := runtime.NumGoroutine()
		require.Equal(t, prevRoutineCnt, currRoutineCnt)

		// because this is "sync" decode, the Write will block and cause a deadlock. We
		// must use another goroutine. This is another sign that the decode is sync and
		// previous synctest is async.
		go func() {
			_, err := pipeWriter.Write(compressedData)
			assert.NoError(t, err)
			err = pipeWriter.Close()
			assert.NoError(t, err)
		}()