Merge pull request #154 from eulerfx/fix_152

ensure offset commit loop reset on rebalance

Author: eulerfx

RELEASE_NOTES.md

@@ -1,3 +1,9 @@
+### 0.1.2-alpha - 03.07.2017
+
+* Fixed consumer offset commit [bug](https://github.com/jet/kafunk/issues/152) wherein after a rebalance a consumer gets assigned a new partition
+  which doesn't receive any messages for longer than the offset retention period, the offsets would be lost.
+  This would only happen after a rebalance not after initial join.
+
 ### 0.1.1-alpha - 25.05.2017
 
 * Snappy compression.

src/kafunk/Consumer.fs

@@ -570,6 +570,10 @@ module Consumer =
     let! state = Group.stateInternal c.groupMember
     return consumerStateFromGroupMemberState state.state }
 
+  /// Returns the stream of consumer states as of the invocation, including the current state.
+  let states (c:Consumer) : AsyncSeq<ConsumerState> =
+    Group.states c.groupMember |> AsyncSeq.map (consumerStateFromGroupMemberState)
+
   let private combineFetchResponses 
     (r1:(ConsumerMessageSet[] * (Partition * HighwaterMarkOffset)[]) option) 
     (r2:(ConsumerMessageSet[] * (Partition * HighwaterMarkOffset)[]) option) =
@@ -862,32 +866,32 @@ module Consumer =
     (c:Consumer)
     (commitInterval:TimeSpan)
     (handler:ConsumerState -> ConsumerMessageSet -> Async<unit>) : Async<unit> = async {
-
       let assignedPartitions = state c |> Async.map (fun s -> s.assignments)
-
-      use commitQueue = Offsets.createPeriodicCommitQueue (commitInterval, assignedPartitions, commitOffsets c)
-            
-      let! currentOffsets = async {
-        let! assignedPartitions = assignedPartitions
-        let! currentOffsets = fetchOffsets c.conn c.config.groupId [| c.config.topic, assignedPartitions |]
-        return
-          currentOffsets 
-          |> Seq.choose (fun (t,os) -> if t = c.config.topic then Some os else None)
-          |> Seq.concat
-          |> Seq.where (fun (_,o) -> o <> -1L)
-          |> Seq.toArray }
-      
-      // commit current offets so that they're in-memory, and will be committed periodically
-      // even if no messages are consumed
-      Offsets.enqueuePeriodicCommit commitQueue currentOffsets
-
+      use commitQueue = Offsets.createPeriodicCommitQueue (commitInterval, assignedPartitions, commitOffsets c)            
+      // commit current offets on group rebalance, including first join,
+      // so that they're committed periodically even if no messages are consumed
+      let commitCurrentOffsetsOnGroupJoinProc =
+        states c
+        |> AsyncSeq.iterAsync (fun s -> async {
+          let! currentOffsets = async {
+            let! currentOffsets = fetchOffsets c.conn c.config.groupId [| c.config.topic, s.assignments |]
+            return
+              currentOffsets 
+              |> Seq.choose (fun (t,os) -> if t = c.config.topic then Some os else None)
+              |> Seq.concat
+              |> Seq.where (fun (_,o) -> o <> -1L)
+              |> Seq.toArray }
+          Offsets.enqueuePeriodicCommit commitQueue currentOffsets })
       let handler s ms = async {
         do! handler s ms
-        Offsets.enqueuePeriodicCommit commitQueue (ConsumerMessageSet.commitPartitionOffsets ms) }
-
-      do! consume c handler
-
-      do! Offsets.flushPeriodicCommit commitQueue }
+        Offsets.enqueuePeriodicCommit commitQueue (ConsumerMessageSet.commitPartitionOffsets ms) }      
+      do!
+        Async.parallel2
+          (async {
+            do! consume c handler
+            do! Offsets.flushPeriodicCommit commitQueue },
+          commitCurrentOffsetsOnGroupJoinProc)
+        |> Async.Ignore }
 
   /// Closes the consumer and leaves the consumer group.
   /// This causes all underlying streams to complete.

src/kafunk/Group.fs

@@ -77,7 +77,7 @@ and GroupConfig = {
 and [<NoEquality;NoComparison;AutoSerializable(false)>] GroupMember = internal {
   conn : KafkaConn
   config : GroupConfig
-  state : MVar<GroupMemberStateWrapper>
+  state : SVar<GroupMemberStateWrapper>
 }
 
 /// The action to take upon leaving the group.
@@ -150,17 +150,23 @@ module Group =
 
   /// Leaves a group, sending a leave group request to Kafka.
   let leave (gm:GroupMember) = async {
-    let! state = MVar.get gm.state
+    let! state = SVar.get gm.state
     return! leaveInternal gm state }
 
   let internal stateInternal (gm:GroupMember) = 
-    gm.state |> MVar.get
+    gm.state |> SVar.get
 
   /// Returns the group member state.
   let state (gm:GroupMember) = async {
     let! state = stateInternal gm
     return state.state }
 
+  /// Returns a stream of group member states as of the invocation, including the current state.
+  let states (gm:GroupMember) =
+    gm.state
+    |> SVar.tap
+    |> AsyncSeq.map (fun s -> s.state)
+
   /// Joins a group.
   let join (gm:GroupMember) (prevState:GroupMemberState option, prevErrorCode:ErrorCode option) = async {
 
@@ -287,7 +293,7 @@ module Group =
       let cts = CancellationTokenSource.CreateLinkedTokenSource (ct, state.state.closed)
       Async.Start (heartbeatProcess, cts.Token)
 
-      let! _ = gm.state |> MVar.put state
+      let! _ = gm.state |> SVar.put state
       return () }
 
 
@@ -360,6 +366,9 @@ module Group =
 
     return! joinSyncHeartbeat (prevState |> Option.map (fun s -> s.memberId), prevErrorCode) }
 
+//  let internal generationsInternal (gm:GroupMember) =
+//    SVar.tap gm.state
+
   let internal generationsInternal (gm:GroupMember) =
     let rec loop () = asyncSeq {
       let! state = stateInternal gm
@@ -385,8 +394,22 @@ module Group =
   let createJoin (conn:KafkaConn) (config:GroupConfig) = async {
     let gm = 
       { GroupMember.conn = conn
-        state = MVar.create ()
+        state = SVar.create ()
         config = config }
+
     let! _ = join gm (None, None)
+    
+//    let rec rejoinProc () = async {
+//      let! state = stateInternal gm
+//      let! action = IVar.get state.closed
+//      match action with
+//      | GroupLeaveAction.LeaveAndRejoin ec ->
+//        let! _ = join gm (Some state.state, Some ec)
+//        return! rejoinProc ()
+//      | GroupLeaveAction.LeaveGroup -> 
+//        return () }
+//    
+//    let! _ = Async.StartChild (rejoinProc ())
+
     return gm }
 

src/kafunk/Utility/SVar.fs

@@ -0,0 +1,76 @@
+[<AutoOpen>]
+module internal Kafunk.SVar
+
+open FSharp.Control
+open System.Threading
+
+/// A stream variable.
+type SVar<'a> = private {
+  state : MVar<'a>
+  tail : SVarNode<'a> ref }
+
+and SVarNode<'a> = {
+  ivar : IVar<('a * SVarNode<'a>)> }
+
+/// Operations on stream variables.
+module SVar =
+  
+  let private createNode () =
+    { ivar = IVar.create () }
+
+  let private putNode (s:SVar<'a>) (a:'a) =
+    let newTail = createNode ()
+    let tail = Interlocked.Exchange(s.tail, newTail)
+    IVar.put (a,newTail) tail.ivar
+
+  /// Creates an empty stream variable.
+  let create () : SVar<'a> = 
+    { state = MVar.create () ; tail = ref (createNode ()) }
+
+  /// Puts a new value into the stream variable.
+  let put (a:'a) (s:SVar<'a>) : Async<unit> = async {
+    let! a = MVar.put a s.state
+    do putNode s a
+    return () }
+
+  /// Atomically updates a value in a stream variable.
+  let updateAsync (f:'a -> Async<'a>) (s:SVar<'a>) : Async<'a> =
+    s.state
+    |> MVar.updateAsync (fun a -> async {
+      let! a' = f a
+      do putNode s a'
+      return a' })
+
+  /// Atomically updates a value in a stream variable.
+  let updateStateAsync (f:'a -> Async<'a * 's>) (s:SVar<'a>) : Async<'s> =
+    s.state
+    |> MVar.updateStateAsync (fun a -> async {
+      let! a',st = f a
+      do putNode s a'
+      return a',st })
+
+  /// Gets the current value of the stream variable.
+  let get (s:SVar<'a>) : Async<'a> =
+    MVar.get s.state
+
+  let rec private tapImpl (s:SVarNode<'a>) : AsyncSeq<'a> = 
+    asyncSeq {
+      let! (a,tl) = IVar.get s.ivar
+      yield a
+      yield! tapImpl tl }
+
+  /// Returns a stream corresponding to changes in the stream variable from the time of invocation, starting with the 
+  /// current value, if any.
+  let tap (s:SVar<'a>) : AsyncSeq<'a> =
+    let curr = MVar.getFastUnsafe s.state    
+    let tail = 
+      let tail = !s.tail
+      match curr with
+      | None -> tail
+      | Some a -> { ivar = IVar.createFull (a,tail) }
+    tapImpl tail
+      
+  /// Publishes an error to all waiting readers of a tapped stream.
+  let error ex (s:SVar<'a>) =
+    let tail = !s.tail
+    IVar.error ex tail.ivar

src/kafunk/kafunk.fsproj

@@ -68,6 +68,7 @@
     <Compile Include="Utility\Crc.fs" />
     <Compile Include="Utility\Async.fs" />
     <Compile Include="Utility\MVar.fs" />
+    <Compile Include="Utility\SVar.fs" />
     <Compile Include="Utility\BoundedMb.fs" />
     <Compile Include="Utility\AsyncSeq.fs" />
     <Compile Include="Utility\Faults.fs" />

tests/kafunk.Tests/SVarTests.fs

@@ -0,0 +1,55 @@
+module SVarTests
+
+open System
+open Kafunk
+open Kafunk.Testing
+open Kafunk.SVar
+open FSharp.Control
+open NUnit.Framework
+
+let asyncSeqToList (s:AsyncSeq<'a>) =
+  let ls = 
+    s
+    |> AsyncSeq.toListAsync 
+    |> Async.timeoutOption (TimeSpan.FromMilliseconds 100.0) 
+    |> Async.RunSynchronously
+  match ls with
+  | Some ls -> ls
+  | None -> failwith "AsyncSeq timed out!"
+
+[<Test>]
+let ``SVar.tap should return stream of all values put after tap`` () =   
+  for put in [ [] ; [1] ; [1;2;3;4] ] do
+    let v = SVar.create ()
+    let stream = SVar.tap v
+    async {
+      for x in put do
+        do! SVar.put x v } |> Async.RunSynchronously
+    let actual = stream |> AsyncSeq.take put.Length |> asyncSeqToList
+    shouldEqual put actual None
+
+[<Test>]
+let ``SVar.tap should return current value`` () = 
+  let v = SVar.create ()
+  let expected = 1
+  SVar.put expected v |> Async.RunSynchronously
+  let expected = [expected]
+  let actual = SVar.tap v |> AsyncSeq.take 1 |> asyncSeqToList
+  shouldEqual expected actual None
+
+[<Test>]
+let ``SVar.tap should return current and subsequent values`` () = 
+  for expected in [ [1] ; [1;2] ; [1;2;3;4;5] ] do
+    let v = SVar.create ()
+    SVar.put 0 v |> Async.RunSynchronously // skipped
+    SVar.put expected.Head v |> Async.RunSynchronously
+    let tapped = SVar.tap v
+    async {
+      for x in expected.Tail do
+        do! SVar.put x v } |> Async.RunSynchronously
+    let actual = tapped |> AsyncSeq.take expected.Length |> asyncSeqToList
+    shouldEqual expected actual None
+
+
+  
+  

tests/kafunk.Tests/kafunk.Tests.fsproj

@@ -76,6 +76,7 @@
     <Compile Include="AsyncTests.fs" />
     <Compile Include="FaultsTests.fs" />
     <Compile Include="AssignmentStrategyTests.fs" />
+    <Compile Include="SVarTests.fs" />
     <None Include="Refs.fsx" />
     <Compile Include="Metadata.fsx" />
     <None Include="Producer.fsx" />