From 598f814f922f9c9523bc6656cd7d1d9043d63726 Mon Sep 17 00:00:00 2001
From: Ryan Wynn <ryan.m.wynn@gmail.com>
Date: Thu, 21 Nov 2019 19:19:53 -0500
Subject: [PATCH] cleanup save state logic

---
 mongofluxd.go | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/mongofluxd.go b/mongofluxd.go
index 1819a57..1170598 100644
--- a/mongofluxd.go
+++ b/mongofluxd.go
@@ -176,8 +176,14 @@ func (im *InfluxMeasure) parseView(view string) error {
 
 func (ctx *InfluxCtx) saveTs() (err error) {
 	if ctx.config.Resume && ctx.lastTs.T > 0 {
+		if err = ctx.writeBatch(); err != nil {
+			return err
+		}
 		if ctx.config.ResumeStrategy == tokenResumeStrategy {
 			err = saveTokens(ctx.client, ctx.tokens, ctx.config)
+			if err == nil {
+				ctx.tokens = bson.M{}
+			}
 		} else {
 			err = saveTimestamp(ctx.client, ctx.lastTs, ctx.config)
 		}
@@ -1043,11 +1049,10 @@ func main() {
 				case op, open := <-gtmCtx.OpC:
 					if op == nil {
 						if !open {
-							if err := influx.writeBatch(); err != nil {
+							if err := influx.saveTs(); err != nil {
 								exitStatus = 1
 								errorLog.Println(err)
 							}
-							influx.saveTs()
 							return
 						}
 						break