Improve performance of Files.walk on the JVM

io/jvm-native/src/main/scala/fs2/io/file/FilesPlatform.scala

@@ -24,6 +24,7 @@ package io
 package file
 
 import cats.effect.kernel.{Async, Resource, Sync}
+import cats.effect.std.Dispatcher
 import cats.syntax.all._
 
 import java.nio.channels.{FileChannel, SeekableByteChannel}
@@ -32,15 +33,16 @@ import java.nio.file.attribute.{
   BasicFileAttributeView,
   BasicFileAttributes => JBasicFileAttributes,
   PosixFileAttributes => JPosixFileAttributes,
-  PosixFilePermissions
+  PosixFilePermissions,
+  FileTime
 }
 import java.security.Principal
 import java.util.stream.{Stream => JStream}
 
 import scala.concurrent.duration._
 
+import fs2.concurrent.Channel
 import fs2.io.CollectionCompat._
-import java.nio.file.attribute.FileTime
 
 private[file] trait FilesPlatform[F[_]] extends DeprecatedFilesApi[F] { self: Files[F] =>
 
@@ -389,6 +391,54 @@ private[file] trait FilesCompanionPlatform {
         .resource(Resource.fromAutoCloseable(javaCollection))
         .flatMap(ds => Stream.fromBlockingIterator[F](collectionIterator(ds), pathStreamChunkSize))
 
+    override def walk(start: Path, maxDepth: Int, followLinks: Boolean): Stream[F, Path] =
+      Stream.resource(Dispatcher.sequential[F]).flatMap { dispatcher =>
+        Stream.eval(Channel.bounded[F, Chunk[Path]](10)).flatMap { channel =>
+          val doWalk = Sync[F].interruptibleMany {
+            val bldr = Vector.newBuilder[Path]
+            val limit = 4096
+            var size = 0
+            bldr.sizeHint(limit)
+            JFiles.walkFileTree(
+              start.toNioPath,
+              if (followLinks) Set(FileVisitOption.FOLLOW_LINKS).asJava else Set.empty.asJava,
+              maxDepth,
+              new SimpleFileVisitor[JPath] {
+                private def enqueue(path: JPath): FileVisitResult = {
+                  bldr += Path.fromNioPath(path)
+                  size += 1
+                  if (size >= limit) {
+                    val result = dispatcher.unsafeRunSync(channel.send(Chunk.from(bldr.result())))
+                    bldr.clear()
+                    size = 0
+                    if (result.isRight) FileVisitResult.CONTINUE else FileVisitResult.TERMINATE
+                  } else FileVisitResult.CONTINUE
+                }
+
+                override def visitFile(file: JPath, attrs: JBasicFileAttributes): FileVisitResult =
+                  enqueue(file)
+
+                override def visitFileFailed(file: JPath, t: IOException): FileVisitResult =
+                  FileVisitResult.CONTINUE
+
+                override def preVisitDirectory(dir: JPath, attrs: JBasicFileAttributes)
+                    : FileVisitResult =
+                  enqueue(dir)
+
+                override def postVisitDirectory(dir: JPath, t: IOException): FileVisitResult =
+                  FileVisitResult.CONTINUE
+              }
+            )
+
+            dispatcher.unsafeRunSync(
+              if (size > 0) channel.closeWithElement(Chunk.from(bldr.result()))
+              else channel.close
+            )
+          }
+          channel.stream.unchunks.concurrently(Stream.eval(doWalk))
+        }
+      }
+
     def createWatcher: Resource[F, Watcher[F]] = Watcher.default(this, F)
 
     def watch(

io/jvm/src/test/scala/fs2/io/file/WalkBenchmark.scala

@@ -1,4 +1,3 @@
-
 /*
  * Copyright (c) 2013 Functional Streams for Scala
  *
@@ -33,7 +32,7 @@ class WalkBenchmark extends Fs2IoSuite {
   private var target: Path = _
 
   override def beforeAll() = {
-		super.beforeAll()
+    super.beforeAll()
     val file = File.createTempFile("fs2-benchmarks-", "-walk")
     file.delete()
     file.mkdir()
@@ -42,21 +41,20 @@ class WalkBenchmark extends Fs2IoSuite {
     val MaxDepth = 7
     val Names = 'A'.to('E').toList.map(_.toString)
 
-    def loop(cwd: File, depth: Int): Unit = {
+    def loop(cwd: File, depth: Int): Unit =
       if (depth < MaxDepth) {
-        Names foreach { name =>
+        Names.foreach { name =>
           val sub = new File(cwd, name)
           sub.mkdir()
           loop(sub, depth + 1)
         }
       } else if (depth == MaxDepth) {
-        Names foreach { name =>
+        Names.foreach { name =>
           val sub = new File(cwd, name)
           sub.createNewFile()
           loop(sub, depth + 1)
         }
       }
-    }
 
     loop(file, 0)
   }
@@ -67,17 +65,21 @@ class WalkBenchmark extends Fs2IoSuite {
     (System.nanoTime() - start).nanos
   }
 
-
-	test("Files.walk has similar performance to java.nio.file.Files.walk") {
-    val fs2Time = time(Files[IO]
-      .walk(target)
-      .compile
-      .count
-      .unsafeRunSync())
+  test("Files.walk has similar performance to java.nio.file.Files.walk") {
+    val fs2Time = time(
+      Files[IO]
+        .walk(target)
+        .compile
+        .count
+        .unsafeRunSync()
+    )
     val nioTime = time(java.nio.file.Files.walk(target.toNioPath).count())
     val epsilon = nioTime.toNanos * 1.5
     println(s"fs2 took: ${fs2Time.toMillis} ms")
     println(s"nio took: ${nioTime.toMillis} ms")
-    assert((fs2Time - nioTime).toNanos.abs < epsilon, s"fs2 time: $fs2Time, nio time: $nioTime, diff: ${fs2Time - nioTime}") 
+    assert(
+      (fs2Time - nioTime).toNanos.abs < epsilon,
+      s"fs2 time: $fs2Time, nio time: $nioTime, diff: ${fs2Time - nioTime}"
+    )
   }
 }