Use LoadErrorHandlingPolicy for Downloads/Offline media segment requests

I am using Exoplayer 2.16.1 in my application. I am trying to control number of segment calls during download for offline playback.

I am trying to control the number of segment API calls in case of HTTP error using custom LoadErrorHandlingPolicy.

This is the high level flow.

val errorHandlingPolicy = CustomDownloadErrorHandlingPolicy(maxRetryCount)

val factory = DashMediaSource.Factory(dataSourceFactory)

val mediaSource = factory
            .setLoadErrorHandlingPolicy(errorHandlingPolicy)
            .createMediaSource(mediaItem)

val downloadHelper = DownloadHelper(
                    mediaItem,
                    mediaSource,
                    DefaultTrackSelector.Parameters.DEFAULT_WITHOUT_CONTEXT,
                    DownloadHelper.getRendererCapabilities(DefaultRenderersFactory(context))

downloadHelper.prepare(this@DownloadPrepareHelper)

I observed that if there is some HTTP error for manifest request, the overridden method of my CustomDownloadErrorHandlingPolicy object is called. (getMinimumLoadableRetryCount(), getRetryDelayMsFor() etc.)

But if there is some error for Segment requests, the callbacks are not getting called. The Exoplayer is trying ~4 times and then throwing error.

With this I am not able to set the right retry count.

Is there a way to set the LoadErrorHandlingPolicy for segments? Or is it a bug in Exoplayer?