Fix propagation of FiberRef modified within a ZQuery (#515)

Author: kyri-petrou

zio-query/shared/src/main/scala/zio/query/ZQuery.scala

@@ -639,17 +639,16 @@ final class ZQuery[-R, +E, +A] private (private val step: ZIO[R, Nothing, Result
             }
             state.setFiberRefs(newRefs)
             restore(runToZIO).exitWith { exit =>
-              val curRefs = state.getFiberRefs(false)
+              var curRefs = state.getFiberRefs(false)
               if (curRefs eq newRefs) {
                 // Cheap and common: FiberRefs were not modified during the execution so we just replace them with the old ones
                 state.setFiberRefs(oldRefs)
               } else {
-                // FiberRefs were mdified so we need to manually revert each one
-                var revertedRefs = oldRefs
-                revertedRefs = resetRef(fid, oldRefs, revertedRefs)(currentCache)
-                revertedRefs = resetRef(fid, oldRefs, revertedRefs)(currentScope)
-                revertedRefs = resetRef(fid, oldRefs, revertedRefs)(disabledCache)
-                state.setFiberRefs(revertedRefs)
+                // FiberRefs were modified so we need to manually revert each one
+                curRefs = resetRef(fid, oldRefs, curRefs)(currentCache)
+                curRefs = resetRef(fid, oldRefs, curRefs)(currentScope)
+                curRefs = resetRef(fid, oldRefs, curRefs)(disabledCache)
+                state.setFiberRefs(curRefs)
               }
               scope.closeAndExitWith(exit)
             }

zio-query/shared/src/test/scala/zio/query/ZQuerySpec.scala

@@ -402,6 +402,18 @@ object ZQuerySpec extends ZIOBaseSpec {
             } yield (c1, c2)
 
           q.run.map { case (c1, c2) => assertTrue(c1 != QueryScope.NoOp, c1 == c2) }
+        },
+        test("propagates FiberRef changes") {
+          val ref = FiberRef.unsafe.make("a")(Unsafe)
+          for {
+            _    <- ZQuery.fromZIO(ref.update(_ + "b")).run
+            res1 <- ref.get
+            _    <- ZQuery.fromZIO(ref.update(_ + "c")).run
+            res2 <- ref.get
+            _    <- ref.set("d")
+            _    <- ZQuery.fromZIO(ref.update(_ + "e")).run
+            res3 <- ref.get
+          } yield assertTrue(res1 == "ab", res2 == "abc", res3 == "de")
         }
       ),
       suite("catchAllZIO")(