FlatMapGroupWithState behaviour when state exists - scala

Question

Spark group state never expires if after a new state is set (with expiration) the group gets "called" for a subsequent time before the initial expiration.
I am setting a timeout of "5 seconds" for any group state. If I send an event which creates a new state and I wait for 5 seconds the group times out successfully. However if I send an event which creates a new state and I send a similar event which also falls in the same group, this group state doesn't time out ever.
// in main class
val aggregation = inputStream
.groupByKey(r => r.key)
.flatMapGroupsWithState(OutputMode.Append, GroupStateTimeout.ProcessingTimeTimeout)(Aggregator.aggregate)
// in another object in another file
object Aggregator {
def aggregate(key: String,
rows: Iterator[InputSchema],
state: GroupState[List[InputSchema]]): Iterator[InputSchema] = {
if (state.hasTimedOut) {
print("time out")
val output = state.get
state.remove()
return output.toIterator
}
if (state.exists) {
println("exists")
state.update(state.get ++ rows.toList)
return Iterator()
}
println("coming to update")
state.update(rows.toList)
state.setTimeoutDuration("5 seconds")
Iterator()
}
}
With the above code, the below happens
1) If i send a message where r.key is 12345,
Print message coming to update gets printed
After 5 seconds time out gets printed.
^ Both are expected behaviour
2) If I send a message where r.key is abcd at time t1, and then I send the same message at t2 when t2-t1 is less than 5 seconds (which is the group timeout initially set), following happens -
Print message coming to update gets printed (this is for the first event at t1)
Print message exists gets printed at t2
Even after waiting for 5 seconds, the state never timesout and I don't get the time out print.
Expected behaviour should be that after 5 seconds from t1 the state should time out.