 (deftype KafkaReadMessages 
   [log-prefix task-map topic ^:unsynchronized-mutable kpartitions batch-timeout
    deserializer-fn segment-fn read-offset ^:unsynchronized-mutable consumer 
-   ^:unsynchronized-mutable iter ^:unsynchronized-mutable offset ^:unsynchronized-mutable drained]
+   ^:unsynchronized-mutable iter ^:unsynchronized-mutable partition->offset ^:unsynchronized-mutable drained]
   p/Plugin
   (start [this event]
     (let [{:keys [kafka/group-id kafka/consumer-opts]} task-map


  p/Checkpointed
  (checkpoint [this]
    offset)

  (recover! [this replica-version checkpoint]
    (set! drained false)
    (set! iter nil)
    (set! offset checkpoint)
    (seek-offset! log-prefix consumer kpartitions task-map topic checkpoint)
    this)


  (poll! [this _]
    (if (and iter (.hasNext ^java.util.Iterator iter))
      (let [rec ^ConsumerRecord (.next ^java.util.Iterator iter)
            deserialized (some-> rec segment-fn)]
        (cond (= :done deserialized)
              (do (set! drained true)
                  nil)
              deserialized
              (let [new-offset (.offset rec)
                    part (.partition rec)]
                (.set ^AtomicLong read-offset new-offset)
                (set! partition->offset (assoc partition->offset part new-offset))
                deserialized)))
      (do (set! iter 
                (.iterator ^ConsumerRecords 
                           (.poll ^Consumer (.consumer ^FranzConsumer consumer) 
