Building ZIO Web App. Part 5. ZIO basic primitives.

If you’ve always been curious about Scala, Functional Programming, and ZIO, then welcome to my series of articles. In this series, I will introduce you to the fundamental principles and approaches of functional programming with Scala and ZIO through the example of building a simple web application.

In this concluding article, we will finalize our application by implementing the remaining components using ZIO’s fundamental primitives.

Articles

1. Building ZIO Web App. Part 1. ZIO and zio-http
2. Building ZIO Web App. Part 2. Unit testing
3. Building ZIO Web App. Part 3. ZLayers and Dependency Injection
4. Building ZIO Web App. Part 4. Designing ZIO application
5. Building ZIO Web App. Part 5. ZIO basic primitives

You can refer to the final version of the source code in my GitHub repository.

To get our application fully operational, we need to implement the last of the low-level services: TokenStorageService and ContentService. To do this, we’ll need to work with a few key ZIO primitives — understanding their purpose, how they function, their behavior in concurrent environments, and of course how to test them.

It’s important to note that I’ll cover only a few of these primitives here. When building a more complex, real-world application, you’ll likely need to leverage a wider range of ZIO tools, so be sure to consult the documentation for guidance tailored to your specific use case.

Ref

Let’s begin with the most straightforward and intuitive approach. Since we’re working with state in a potentially concurrent environment, Ref appears to be a solid choice.

Our goal is to model a storage system for our tokens. In a world without ZIO and effects, the simplest approach would be to use a HashMap where the userId serves as the key and the Token as the value. However, because this represents mutable state, we will wrap our HashMap in a ZIO Ref type. This will give us access to a suite of methods for updating, mutating, and retrieving values, all of which are implemented as ZIO effects.

To create a Ref with a Map, we can use the Ref.make method. Since this method also returns an effect, we can begin our implementation as follows:

object TokenStorageService:
  def apply(): ZIO[Any, Nothing, TokenStorageService] = for {
    mapRef <- Ref.make(Map.empty[Long, Token])
  } yield new TokenStorageService:
//...

Now that we have our mapRef, we can use it in our implementation:

extension (r: Ref[Map[Long, Token]])
  def generateAndUpdateToken(userId: Long): ZIO[Any, Nothing, Token] = for {
    uuid <- Random.nextUUID
    time <- Clock.currentDateTime
    token = Token(uuid.toString, time.plusHours(1L))
    _    <- r.update(m => m + (userId -> token))
  } yield token

override def getOrCreateToken(userId: Long): ZIO[Any, Nothing, Token] =
  for {
    map         <- mapRef.get
    oldToken    =  map.get(userId)
    actualToken <- oldToken match {
      case Some(token) => for {
        currentTime  <- Clock.currentDateTime
        updatedToken <- if (token.expires.isBefore(currentTime)) mapRef.generateAndUpdateToken(userId)
                        else ZIO.succeed(token)
      } yield updatedToken
      case None => mapRef.generateAndUpdateToken(userId)
    }
  } yield actualToken

override def checkExpired(userId: Long, token: _root_.java.lang.String): ZIO[Any, UserNotFoundError | AccessDenied, Boolean] = ???

• To keep the code clean and concise, we’ve extracted the generateAndUpdateToken method as an extension for our mapRef.
• The checkExpired method will be implemented later
• In the getOrCreateToken method, we first fetch the map and then retrieve the token associated with the user.
• Since the token might not exist for the user, we use pattern matching to check its presence. If the token is missing or has expired, we generate a new one using the generateAndUpdateToken extension method. Otherwise, we simply return the existing token.

Notice that we are primarily working with the ZIO type: Random.nextUUID, Clock.currentDateTime, mapRef.get, and others—all of which return effects. This allows us to leverage the full power of ZIO along with Scala's type system and syntactic sugar, such as for-comprehensions.

For those experienced with concurrency, a potential issue might be evident: if multiple requests are submitted simultaneously, a token could be issued more than once. Before addressing this issue, it’s beneficial to first demonstrate the potential problem with a test. This approach will not only highlight the bug but also help prevent it from reoccurring in the future.

Test

With the ZIO.foreachPar method, we can easily simulate any number of concurrent requests to our service. Let’s proceed with this approach and create src/test/scala/TokenStorageServiceSute.scala file:

import zio._
import zio.test._

object TokenStorageServiceSuite extends ZIOSpecDefault:
  def spec = suite("getOrCreateToken") {
    test("Concurrent access") {
      for {
        tokenStorageService <- ZIO.service[TokenStorageService]
        list <- ZIO.foreachPar(1 to 100)(_ => tokenStorageService.getOrCreateToken(1L))
      } yield assertTrue(list.distinct.size == 1)
    } @@ TestAspect.repeat(Schedule.recurs(10))
  }.provide(ZLayer.fromZIO(TokenStorageService.apply()))

• Using ZIO.foreachPar, we execute 100 requests for getOrCreateToken in parallel and collect the results in a list.
• The correct behavior should be to return only one unique token. We verify this by checking the number of distinct elements in the list.
• Since the test is not deterministic, there is a theoretical chance that it might pass occasionally by chance. To minimize this likelihood, we use TestAspect.repeat to run the test 10 times, increasing our confidence in the correctness of the implementation.

Note that we need to provide a ZLayer with an implementation of TokenStorageService. Since creating the service with TokenStorageService.apply() results in an effect itself, we can use the ZLayer.fromZIO method to provide the layer.

Now, we can run the test and clearly observe that it is failing!

Semaphore

To address this issue, the first approach that comes to mind is to use a Semaphore. A Semaphore is a concurrency control structure that can limit parallel operations by providing locks and releases. We can use it to wrap our manipulations with mapRef using the withPermit method. This ensures that only one operation can modify the mapRef at a time, preventing multiple tokens from being issued simultaneously:

object TokenStorageService:
  def apply(): ZIO[Any, Nothing, TokenStorageService] = for {
    mapRef    <- Ref.make(Map.empty[Long, Token])
    semaphore <- Semaphore.make(1)
  } yield new TokenStorageService:
    extension (r: Ref[Map[Long, Token]])
      def generateAndUpdateToken(userId: Long): ZIO[Any, Nothing, Token] = for {
        uuid <- Random.nextUUID
        time <- Clock.currentDateTime
        token = Token(uuid.toString, time.plusHours(1L))
        _ <- r.update(m => m + (userId -> token))
      } yield token

    override def getOrCreateToken(userId: Long): ZIO[Any, Nothing, Token] =
      semaphore.withPermit(getOrCreateTokenInternal(userId))

    private def getOrCreateTokenInternal(userId: Long): ZIO[Any, Nothing, Token] =
      for {
        map         <- mapRef.get
        oldToken    =  map.get(userId)
        actualToken <- oldToken match {
          case Some(token) => for {
            currentTime  <- Clock.currentDateTime
            updatedToken <- if (token.expires.isBefore(currentTime)) mapRef.generateAndUpdateToken(userId)
            else ZIO.succeed(token)
          } yield updatedToken
          case None => mapRef.generateAndUpdateToken(userId)
        }
      } yield actualToken

    override def checkExpired(userId: Long, token: _root_.java.lang.String): ZIO[Any, UserNotFoundError | AccessDenied, Boolean] = ???

The only change here is that we added a Semaphore to guard our code. We initialized it at the very top by passing the number of permits (1) using Semaphore.make(1), which ensures that only one thread can enter at a time. We then wrapped the existing getOrCreateToken implementation with semaphore.withPermit.

With this adjustment, when we run the test again, it passes. That was a straightforward fix!

Although even our code is working correctly, a bit more experienced programmers will spot an issue even here. And that is performance: for every request we are locking the whole map and not allowing to do anything with it even for other users. Since the tutorial is very basic, we will not dive into it here. Drop in comments if you would like to see the solutions for this problem in the next article.

More tests

Now that we’ve addressed the concurrency issue, we can focus on additional tests: verifying that a token is issued correctly by a request and that the token expires after one hour. While there are many other potential tests we could consider (and the more tests, the better), let’s start with these two.

We can add these tests to the existing test suite using the + operator to combine them:

test("the same token for the same user") {
    for {
      tokenStorageService <- ZIO.service[TokenStorageService]
      token1              <- tokenStorageService.getOrCreateToken(1)
      token2              <- tokenStorageService.getOrCreateToken(1)
    } yield assertTrue(token1 == token2)
  } 

Here, we are verifying that a token cannot be issued twice for the same user. A more compelling test is to check if the token expires as expected. However, waiting for two hours in the test isn’t practical. How can we wait for 2 hours in the test?

test("Token expiration") {
    for {
      tokenStorageService <- ZIO.service[TokenStorageService]
      token1              <- tokenStorageService.getOrCreateToken(1)
      _                   <- TestClock.adjust(Duration.apply(2, TimeUnit.HOURS))
      token2              <- tokenStorageService.getOrCreateToken(1)
    } yield assertTrue(token2 != token1) && assertTrue(token2.expires.isAfter(token1.expires))
  }

The same as for TestRandom ZIO have a tweaks for the Clock. Here with TestClock we’re emulating waiting for 2 hours between token requests. Tokens should be different and the second one obviously should expire after the first one.

checkExpired

Now we can return to the TokenStorageService and implement the checkExpired method:

override def checkExpired(userId: Long, token: String): ZIO[Any, UserNotFoundError | AccessDenied, Boolean] = for {
  map <- mapRef.get
  storedToken = map.get(userId)
  now <- Clock.currentDateTime
  result <- storedToken match {
    case Some(t) if t.uid == token => ZIO.succeed(t.expires.isBefore(now))
    case Some(_) => ZIO.fail(AccessDenied())
    case None => ZIO.fail(UserNotFoundError(userId))
  }
} yield result

Again, Scala’s strong type system aids us with types and implementation. By using Option[Token] for the token type, it encourages us to handle the None case, which allows us to decide to fail the call with a UserNotFound or AccessDenied error if needed.

Even a few lines of code can introduce various corner cases, so it’s crucial to create unit tests to cover them. For now, let’s focus on the first two scenarios: checking the token’s expiration and handling the case when a user is not found. We’ll create a separate test suite for these cases and append it to the spec by +, using the same ZLayer for dependency injection:

suite("checkExpired") {
  test("check expiration") {
    val v  = for {
      tokenStorageService <- ZIO.service[TokenStorageService]
      token               <- tokenStorageService.getOrCreateToken(1)
      _                   <- TestClock.adjust(Duration.apply(2, TimeUnit.HOURS))
      result              <- tokenStorageService.checkExpired(1, token.uid)
    } yield assertTrue(result)
    v
  } +
  test("check user not found") {
    val expired = for {
      tokenStorageService <- ZIO.service[TokenStorageService]
      token               <- tokenStorageService.getOrCreateToken(1)
      result              <- tokenStorageService.checkExpired(2, token.uid)
    } yield result
    assertZIO(expired.exit)(fails(equalTo(UserNotFoundError(2))))
  }
}.provide(ZLayer.fromZIO(TokenStorageService()))

Note that in the second test, we checked the failure case of the ZIO type.

ContentService

Last piece of puzzle is ContentService. To keep things easy, let’s implement the bussiness logic in the most straighforward way: for odd contentIds we will return a string Content <contentId> for others we will fail with NoContentError.

object ContentServiceLive extends ContentService:
  override def accessContent(contentId: Long): ZIO[Any, NoContentError, Content] =
    if(contentId % 2 == 0) ZIO.succeed(s"Content $contentId")
    else ZIO.fail(NoContentError(contentId))

This implementation resembles a mock setup, so I’ll leave unit testing as an exercise for you.

Finally, don’t forget to add our implementations to the list of production ZLayers. This should be done in Main.scala:

  def run = Server.install(apps)
    .flatMap { port =>
      Console.printLine(s"Started on $port") *> ZIO.never
    }
    .provide(Server.defaultWithPort(8888),
      ZLayer.succeed(HealthServiceLive),
      ZLayer.succeed(BusinessServiceLive),
      ZLayer.fromZIO(TokenStorageService()),
      ZLayer.succeed(ContentServiceLive)
    )

Note that to create a ZLayer from ZIO effects, we use the method ZLayer.fromZIO. This is applicable for TokenStorageService, as its apply method returns an effect of type ZIO[Any, Nothing, TokenStorageService].

All together!

We are now ready to run the application. Start it from your IDE or by using the sbt run command, and your endpoints will be available on port 8888. You can use curl to test the behavior of your application. For example, you can check how invalid user IDs are handled and verify that the error messages match the expected responses from exceptions.

$ curl -i -X POST 'localhost:8888/issueToken?userId=asd'
HTTP/1.1 400 Bad Request
warning: 400 ZIO HTTP For input string: "asd"
content-length: 0

Token can be issued, but it can’t be issued twice times in a row:

$ curl  -X POST 'localhost:8888/issueToken?userId=1' 
{"uid":"f5f8f34e-73e1-410d-b1e1-b7bee0b2d9fd","expires":"2024-04-18T11:33:17.55488+04:00"}

$ curl  -X POST 'localhost:8888/issueToken?userId=1' 
{"uid":"f5f8f34e-73e1-410d-b1e1-b7bee0b2d9fd","expires":"2024-04-18T11:33:17.55488+04:00"}

Content:

# Content is not avalable with invalid token
$ curl -i 'localhost:8888/content?userId=1&token=invalid-token&contentId=2'
HTTP/1.1 403 Forbidden
warning: 403 ZIO HTTP Token is not valid
content-length: 0

$ curl 'localhost:8888/content?userId=1&token=f5f8f34e-73e1-410d-b1e1-b7bee0b2d9fd&contentId=2' 
Content 2

$ curl -i 'localhost:8888/content?userId=1&token=f5f8f34e-73e1-410d-b1e1-b7bee0b2d9fd&contentId=1'
HTTP/1.1 404 Not Found
warning: 404 ZIO HTTP /content
content-length: 0

Conclusion

Congratulations! We’ve just completed our first ZIO web application. Over the course of these five articles, we explored fundamental concepts of functional programming and ZIO, observed how the compiler assists in design and modeling, and helps prevent a variety of errors. We also emphasized the importance of unit testing and successfully integrated all components into a functioning web server.

In these five articles, we have only scratched the surface of the basic concepts of ZIO and functional programming, which allowed us to build a simple web application. If you want to dive deeper into this topic, it’s essential to explore the documentation and other resources in more detail.

For any questions, feel free to contact me via Linkedin or Telegram (tg: @AnKochem).