lock s.kill, send done signal

query/gremlin/session.go

@@ -93,8 +93,11 @@ func (s *Session) SendResult(r *Result) bool {
 	if s.limit >= 0 && s.limit == s.count {
 		return false
 	}
+	s.envLock.Lock()
+	kill := s.kill
+	s.envLock.Unlock()
 	select {
-	case <-s.kill:
+	case <-kill:
 		return false
 	default:
 	}
@@ -125,28 +128,40 @@ func (s *Session) runUnsafe(input interface{}) (otto.Value, error) {
 	s.env.Interrupt = make(chan func(), 1)
 
 	ready := make(chan struct{})
-	s.kill = make(chan struct{})
+	done := make(chan struct{})
 	if s.timeout >= 0 {
 		go func() {
 			time.Sleep(s.timeout)
 			<-ready
-			close(s.kill)
-			s.envLock.Lock()
-			defer s.envLock.Unlock()
-			if s.env != nil {
-				s.env.Interrupt <- func() {
-					panic(ErrKillTimeout)
+			select {
+			case <-done:
+				return
+			default:
+				close(s.kill)
+				s.envLock.Lock()
+				defer s.envLock.Unlock()
+				s.kill = nil
+				if s.env != nil {
+					s.env.Interrupt <- func() {
+						panic(ErrKillTimeout)
+					}
+					s.env = s.emptyEnv
 				}
-				s.env = s.emptyEnv
+				return
 			}
 		}()
 	}
 
 	s.envLock.Lock()
 	env := s.env
+	if s.kill == nil {
+		s.kill = make(chan struct{})
+	}
 	s.envLock.Unlock()
 	close(ready)
-	return env.Run(input)
+	out, err := env.Run(input)
+	close(done)
+	return out, err
 }
 
 func (s *Session) ExecInput(input string, out chan interface{}, limit int) {
@@ -257,8 +272,11 @@ func (s *Session) GetJson() ([]interface{}, error) {
 	if s.err != nil {
 		return nil, s.err
 	}
+	s.envLock.Lock()
+	kill := s.kill
+	s.envLock.Unlock()
 	select {
-	case <-s.kill:
+	case <-kill:
 		return nil, ErrKillTimeout
 	default:
 		return s.dataOutput, nil