From 9549d98a7caeaf04a09e64bb16fd057ab8eb4fee Mon Sep 17 00:00:00 2001
From: miraclesu <suchuangji@gmail.com>
Date: Tue, 26 Mar 2013 17:22:07 +0800
Subject: [PATCH] Close friendly

---
 worker/agent.go  | 6 ++++--
 worker/worker.go | 8 +++++---
 2 files changed, 9 insertions(+), 5 deletions(-)

diff --git a/worker/agent.go b/worker/agent.go
index 9c8cc01..1e17584 100644
--- a/worker/agent.go
+++ b/worker/agent.go
@@ -41,7 +41,7 @@ func newAgent(addr string, worker *Worker) (a *agent, err error) {
 func (a *agent) outLoop() {
     ok := true
     var job *Job
-    for ok {
+    for a.worker.running && ok {
         if job, ok = <-a.out; ok {
             if err := a.write(job.Encode()); err != nil {
                 a.worker.err(err)
@@ -92,7 +92,9 @@ func (a *agent) inLoop() {
             a.WriteJob(newJob(common.REQ, common.GRAB_JOB_UNIQ, nil))
         case common.ERROR, common.ECHO_RES, common.JOB_ASSIGN_UNIQ, common.JOB_ASSIGN:
             job.agent = a
-            a.worker.in <- job
+            if a.worker.running {
+                a.worker.in <- job
+            }
         }
     }
 }
diff --git a/worker/worker.go b/worker/worker.go
index 33f6140..5f26b46 100644
--- a/worker/worker.go
+++ b/worker/worker.go
@@ -165,7 +165,7 @@ func (worker *Worker) removeFunc(funcname string) {
 func (worker *Worker) dealJob(job *Job) {
     defer func() {
         job.Close()
-        if worker.limit != nil {
+        if worker.running && worker.limit != nil {
             worker.limit <- true
         }
     }()
@@ -185,7 +185,6 @@ func (worker *Worker) dealJob(job *Job) {
 // Main loop
 func (worker *Worker) Work() {
     defer func() {
-        worker.running = false
         for _, v := range worker.agents {
             v.Close()
         }
@@ -220,6 +219,7 @@ func (worker *Worker) handleJob(job *Job) {
 
 // Close.
 func (worker *Worker) Close() {
+    worker.running = false
     close(worker.in)
     if worker.limit != nil {
         close(worker.limit)
@@ -299,7 +299,9 @@ func (worker *Worker) exec(job *Job) (err error) {
     job.magicCode = common.REQ
     job.DataType = datatype
     job.Data = r.data
-    job.agent.WriteJob(job)
+    if worker.running {
+        job.agent.WriteJob(job)
+    }
     return
 }