#include <linux/idr.h>
 #include <linux/kthread.h>
 #include <../drivers/ata/ahci.h>
+#include <linux/export.h>
 #include "mtip32xx.h"
 
 #define HW_CMD_SLOT_SZ         (MTIP_MAX_COMMAND_SLOTS * 32)
 #define HW_PORT_PRIV_DMA_SZ \
                (HW_CMD_SLOT_SZ + HW_CMD_TBL_AR_SZ + AHCI_RX_FIS_SZ)
 
+#define HOST_CAP_NZDMA         (1 << 19)
 #define HOST_HSORG             0xFC
 #define HSORG_DISABLE_SLOTGRP_INTR (1<<24)
 #define HSORG_DISABLE_SLOTGRP_PXIS (1<<16)
        int group = 0, commandslot = 0, commandindex = 0;
        struct mtip_cmd *command;
        struct mtip_port *port = dd->port;
+       static int in_progress;
+
+       if (in_progress)
+               return;
+
+       in_progress = 1;
 
        for (group = 0; group < 4; group++) {
                for (commandslot = 0; commandslot < 32; commandslot++) {
 
        up(&port->cmd_slot);
 
-       atomic_set(&dd->drv_cleanup_done, true);
+       set_bit(MTIP_DD_FLAG_CLEANUP_BIT, &dd->dd_flag);
+       in_progress = 0;
 }
 
 /*
                 && time_before(jiffies, timeout))
                mdelay(1);
 
+       if (test_bit(MTIP_DD_FLAG_REMOVE_PENDING_BIT, &dd->dd_flag))
+               return -1;
+
        if (readl(dd->mmio + HOST_CTL) & HOST_RESET)
                return -1;
 
                 && time_before(jiffies, timeout))
                ;
 
+       if (test_bit(MTIP_DD_FLAG_REMOVE_PENDING_BIT, &port->dd->dd_flag))
+               return;
+
        /*
         * Chip quirk: escalate to hba reset if
         * PxCMD.CR not clear after 500 ms
        while (time_before(jiffies, timeout))
                ;
 
+       if (test_bit(MTIP_DD_FLAG_REMOVE_PENDING_BIT, &port->dd->dd_flag))
+               return;
+
        /* Clear PxSCTL.DET */
        writel(readl(port->mmio + PORT_SCR_CTL) & ~1,
                         port->mmio + PORT_SCR_CTL);
                         && time_before(jiffies, timeout))
                ;
 
+       if (test_bit(MTIP_DD_FLAG_REMOVE_PENDING_BIT, &port->dd->dd_flag))
+               return;
+
        if ((readl(port->mmio + PORT_SCR_STAT) & 0x01) == 0)
                dev_warn(&port->dd->pdev->dev,
                        "COM reset failed\n");
        if (unlikely(!port))
                return;
 
-       if (atomic_read(&port->dd->resumeflag) == true) {
+       if (test_bit(MTIP_DD_FLAG_RESUME_BIT, &port->dd->dd_flag)) {
                mod_timer(&port->cmd_timer,
                        jiffies + msecs_to_jiffies(30000));
                return;
                                /* don't proceed further */
                                return IRQ_HANDLED;
                        }
+                       if (test_bit(MTIP_DD_FLAG_REMOVE_PENDING_BIT,
+                                                       &dd->dd_flag))
+                               return rv;
 
                        mtip_process_errors(dd, port_stat & PORT_IRQ_ERR);
                }
                        msleep(20);
                        continue; /* svc thd is actively issuing commands */
                }
+               if (test_bit(MTIP_DD_FLAG_REMOVE_PENDING_BIT,
+                                               &port->dd->dd_flag))
+                       return -EFAULT;
                /*
                 * Ignore s_active bit 0 of array element 0.
                 * This bit will always be set
                                "Internal command did not complete [%d] "
                                "within timeout of  %lu ms\n",
                                atomic, timeout);
+                       if (mtip_check_surprise_removal(port->dd->pdev) ||
+                               test_bit(MTIP_DD_FLAG_REMOVE_PENDING_BIT,
+                                               &port->dd->dd_flag)) {
+                               rv = -ENXIO;
+                               goto exec_ic_exit;
+                       }
                        rv = -EAGAIN;
                }
 
                        & (1 << MTIP_TAG_INTERNAL)) {
                        dev_warn(&port->dd->pdev->dev,
                                "Retiring internal command but CI is 1.\n");
+                       if (test_bit(MTIP_DD_FLAG_REMOVE_PENDING_BIT,
+                                               &port->dd->dd_flag)) {
+                               hba_reset_nosleep(port->dd);
+                               rv = -ENXIO;
+                       } else {
+                               mtip_restart_port(port);
+                               rv = -EAGAIN;
+                       }
+                       goto exec_ic_exit;
                }
 
        } else {
                while ((readl(
                        port->cmd_issue[MTIP_TAG_INTERNAL])
                        & (1 << MTIP_TAG_INTERNAL))
-                       && time_before(jiffies, timeout))
-                       ;
+                       && time_before(jiffies, timeout)) {
+                       if (mtip_check_surprise_removal(port->dd->pdev) ||
+                               test_bit(MTIP_DD_FLAG_REMOVE_PENDING_BIT,
+                                               &port->dd->dd_flag)) {
+                               rv = -ENXIO;
+                               goto exec_ic_exit;
+                       }
+               }
 
                if (readl(port->cmd_issue[MTIP_TAG_INTERNAL])
                        & (1 << MTIP_TAG_INTERNAL)) {
                                "Internal command did not complete [%d]\n",
                                atomic);
                        rv = -EAGAIN;
+                       if (test_bit(MTIP_DD_FLAG_REMOVE_PENDING_BIT,
+                                               &port->dd->dd_flag)) {
+                               hba_reset_nosleep(port->dd);
+                               rv = -ENXIO;
+                       } else {
+                               mtip_restart_port(port);
+                               rv = -EAGAIN;
+                       }
                }
        }
-
+exec_ic_exit:
        /* Clear the allocated and active bits for the internal command. */
        atomic_set(&int_cmd->active, 0);
        release_slot(port, MTIP_TAG_INTERNAL);
        int rv = 0;
        struct host_to_dev_fis fis;
 
+       if (test_bit(MTIP_DD_FLAG_REMOVE_PENDING_BIT, &port->dd->dd_flag))
+               return -EFAULT;
+
        /* Build the FIS. */
        memset(&fis, 0, sizeof(struct host_to_dev_fis));
        fis.type        = 0x27;
        fis.device      = command[6] & ~0x10; /* Clear the dev bit*/
 
 
-       dbg_printk(MTIP_DRV_NAME "%s: User Command: cmd %x, feat %x, "
-               "nsect %x, sect %x, lcyl %x, "
-               "hcyl %x, sel %x\n",
+       dbg_printk(MTIP_DRV_NAME "%s: User Command: cmd %x, feat %x, nsect %x, sect %x, lcyl %x, hcyl %x, sel %x\n",
                __func__,
                command[0],
                command[1],
        command[4] = reply->cyl_low;
        command[5] = reply->cyl_hi;
 
-       dbg_printk(MTIP_DRV_NAME "%s: Completion Status: stat %x, "
-               "err %x , cyl_lo %x cyl_hi %x\n",
+       dbg_printk(MTIP_DRV_NAME "%s: Completion Status: stat %x, err %x , cyl_lo %x cyl_hi %x\n",
                __func__,
                command[0],
                command[1],
        struct host_to_dev_fis  *fis;
        struct mtip_port *port = dd->port;
        struct mtip_cmd *command = &port->commands[tag];
+       int dma_dir = (dir == READ) ? DMA_FROM_DEVICE : DMA_TO_DEVICE;
 
        /* Map the scatter list for DMA access */
-       if (dir == READ)
-               nents = dma_map_sg(&dd->pdev->dev, command->sg,
-                                       nents, DMA_FROM_DEVICE);
-       else
-               nents = dma_map_sg(&dd->pdev->dev, command->sg,
-                                       nents, DMA_TO_DEVICE);
+       nents = dma_map_sg(&dd->pdev->dev, command->sg, nents, dma_dir);
 
        command->scatter_ents = nents;
 
         */
        command->comp_data = dd;
        command->comp_func = mtip_async_complete;
-       command->direction = (dir == READ ? DMA_FROM_DEVICE : DMA_TO_DEVICE);
+       command->direction = dma_dir;
 
        /*
         * Set the completion function and data for the command passed
        down(&dd->port->cmd_slot);
        *tag = get_slot(dd->port);
 
+       if (unlikely(test_bit(MTIP_DD_FLAG_REMOVE_PENDING_BIT, &dd->dd_flag))) {
+               up(&dd->port->cmd_slot);
+               return NULL;
+       }
        if (unlikely(*tag < 0))
                return NULL;
 
  * return value
  *     The size, in bytes, of the data copied into buf.
  */
-static ssize_t hw_show_registers(struct device *dev,
+static ssize_t mtip_hw_show_registers(struct device *dev,
                                struct device_attribute *attr,
                                char *buf)
 {
 
        return size;
 }
-static DEVICE_ATTR(registers, S_IRUGO, hw_show_registers, NULL);
+static DEVICE_ATTR(registers, S_IRUGO, mtip_hw_show_registers, NULL);
 
 /*
  * Create the sysfs related attributes.
                "FTL rebuild in progress. Polling for completion.\n");
 
        start = jiffies;
-       dd->ftlrebuildflag = 1;
        timeout = jiffies + msecs_to_jiffies(MTIP_FTL_REBUILD_TIMEOUT_MS);
 
        do {
+               if (unlikely(test_bit(MTIP_DD_FLAG_REMOVE_PENDING_BIT,
+                               &dd->dd_flag)))
+                       return -EFAULT;
                if (mtip_check_surprise_removal(dd->pdev))
                        return -EFAULT;
 
                        dev_warn(&dd->pdev->dev,
                                "FTL rebuild complete (%d secs).\n",
                        jiffies_to_msecs(jiffies - start) / 1000);
-                       dd->ftlrebuildflag = 0;
                        mtip_block_initialize(dd);
-                       break;
+                       return 0;
                }
                ssleep(10);
        } while (time_before(jiffies, timeout));
 
        /* Check for timeout */
-       if (dd->ftlrebuildflag) {
-               dev_err(&dd->pdev->dev,
+       dev_err(&dd->pdev->dev,
                "Timed out waiting for FTL rebuild to complete (%d secs).\n",
                jiffies_to_msecs(jiffies - start) / 1000);
-               return -EFAULT;
-       }
-
-       return 0;
+       return -EFAULT;
 }
 
 /*
                if (kthread_should_stop())
                        break;
 
+               if (unlikely(test_bit(MTIP_DD_FLAG_REMOVE_PENDING_BIT,
+                               &dd->dd_flag)))
+                       break;
                set_bit(MTIP_FLAG_SVC_THD_ACTIVE_BIT, &port->flags);
                if (test_bit(MTIP_FLAG_ISSUE_CMDS_BIT, &port->flags)) {
                        slot = 1;
        int i;
        int rv;
        unsigned int num_command_slots;
+       unsigned long timeout, timetaken;
 
        dd->mmio = pcim_iomap_table(dd->pdev)[MTIP_ABAR];
 
                        dd->port->mmio + i*0x80 + PORT_SDBV;
        }
 
-       /* Reset the HBA. */
-       if (mtip_hba_reset(dd) < 0) {
-               dev_err(&dd->pdev->dev,
-                       "Card did not reset within timeout\n");
-               rv = -EIO;
+       timetaken = jiffies;
+       timeout = jiffies + msecs_to_jiffies(30000);
+       while (((readl(dd->port->mmio + PORT_SCR_STAT) & 0x0F) != 0x03) &&
+                time_before(jiffies, timeout)) {
+               mdelay(100);
+       }
+       if (unlikely(mtip_check_surprise_removal(dd->pdev))) {
+               timetaken = jiffies - timetaken;
+               dev_warn(&dd->pdev->dev,
+                       "Surprise removal detected at %u ms\n",
+                       jiffies_to_msecs(timetaken));
+               rv = -ENODEV;
+               goto out2 ;
+       }
+       if (unlikely(test_bit(MTIP_DD_FLAG_REMOVE_PENDING_BIT, &dd->dd_flag))) {
+               timetaken = jiffies - timetaken;
+               dev_warn(&dd->pdev->dev,
+                       "Removal detected at %u ms\n",
+                       jiffies_to_msecs(timetaken));
+               rv = -EFAULT;
                goto out2;
        }
 
+       /* Conditionally reset the HBA. */
+       if (!(readl(dd->mmio + HOST_CAP) & HOST_CAP_NZDMA)) {
+               if (mtip_hba_reset(dd) < 0) {
+                       dev_err(&dd->pdev->dev,
+                               "Card did not reset within timeout\n");
+                       rv = -EIO;
+                       goto out2;
+               }
+       } else {
+               /* Clear any pending interrupts on the HBA */
+               writel(readl(dd->mmio + HOST_IRQ_STAT),
+                       dd->mmio + HOST_IRQ_STAT);
+       }
+
        mtip_init_port(dd->port);
        mtip_start_port(dd->port);
 
        mod_timer(&dd->port->cmd_timer,
                jiffies + msecs_to_jiffies(MTIP_TIMEOUT_CHECK_PERIOD));
 
+
+       if (test_bit(MTIP_DD_FLAG_REMOVE_PENDING_BIT, &dd->dd_flag)) {
+               rv = -EFAULT;
+               goto out3;
+       }
+
        if (mtip_get_identify(dd->port, NULL) < 0) {
                rv = -EFAULT;
                goto out3;
         * Send standby immediate (E0h) to the drive so that it
         * saves its state.
         */
-       if (atomic_read(&dd->drv_cleanup_done) != true) {
+       if (!test_bit(MTIP_DD_FLAG_CLEANUP_BIT, &dd->dd_flag)) {
 
-               mtip_standby_immediate(dd->port);
+               if (test_bit(MTIP_FLAG_REBUILD_BIT, &dd->dd_flag))
+                       if (mtip_standby_immediate(dd->port))
+                               dev_warn(&dd->pdev->dev,
+                                       "STANDBY IMMEDIATE failed\n");
 
                /* de-initialize the port. */
                mtip_deinit_port(dd->port);
        if (!dd)
                return -ENOTTY;
 
+       if (unlikely(test_bit(MTIP_DD_FLAG_REMOVE_PENDING_BIT, &dd->dd_flag)))
+               return -ENOTTY;
+
        switch (cmd) {
        case BLKFLSBUF:
                return -ENOTTY;
        if (!dd)
                return -ENOTTY;
 
+       if (unlikely(test_bit(MTIP_DD_FLAG_REMOVE_PENDING_BIT, &dd->dd_flag)))
+               return -ENOTTY;
+
        switch (cmd) {
        case BLKFLSBUF:
                return -ENOTTY;
        int nents = 0;
        int tag = 0;
 
+       if (test_bit(MTIP_DD_FLAG_REMOVE_PENDING_BIT, &dd->dd_flag)) {
+               bio_endio(bio, -ENXIO);
+               return;
+       }
+
        if (unlikely(!bio_has_data(bio))) {
                blk_queue_flush(queue, 0);
                bio_endio(bio, 0);
 
                if (unlikely((bio)->bi_vcnt > MTIP_MAX_SG)) {
                        dev_warn(&dd->pdev->dev,
-                               "Maximum number of SGL entries exceeded");
+                               "Maximum number of SGL entries exceeded\n");
                        bio_io_error(bio);
                        mtip_hw_release_scatterlist(dd, tag);
                        return;
                kobject_put(kobj);
        }
 
-       if (dd->mtip_svc_handler)
+       if (dd->mtip_svc_handler) {
+               set_bit(MTIP_DD_FLAG_INIT_DONE_BIT, &dd->dd_flag);
                return rv; /* service thread created for handling rebuild */
+       }
 
 start_service_thread:
        sprintf(thd_name, "mtip_svc_thd_%02d", index);
                goto kthread_run_error;
        }
 
+       if (wait_for_rebuild == MTIP_FTL_REBUILD_MAGIC)
+               rv = wait_for_rebuild;
+
        return rv;
 
 kthread_run_error:
        }
 
        /* Clean up the sysfs attributes managed by the protocol layer. */
-       kobj = kobject_get(&disk_to_dev(dd->disk)->kobj);
-       if (kobj) {
-               mtip_hw_sysfs_exit(dd, kobj);
-               kobject_put(kobj);
+       if (test_bit(MTIP_DD_FLAG_INIT_DONE_BIT, &dd->dd_flag)) {
+               kobj = kobject_get(&disk_to_dev(dd->disk)->kobj);
+               if (kobj) {
+                       mtip_hw_sysfs_exit(dd, kobj);
+                       kobject_put(kobj);
+               }
        }
 
        /*
                return -ENOMEM;
        }
 
-       atomic_set(&dd->resumeflag, false);
-
        /* Attach the private data to this PCI device.  */
        pci_set_drvdata(pdev, dd);
 
         * instance number.
         */
        instance++;
-
+       if (rv != MTIP_FTL_REBUILD_MAGIC)
+               set_bit(MTIP_DD_FLAG_INIT_DONE_BIT, &dd->dd_flag);
        goto done;
 
 block_initialize_err:
        pci_set_drvdata(pdev, NULL);
        return rv;
 done:
-       /* Set the atomic variable as 0 */
-       atomic_set(&dd->drv_cleanup_done, false);
-
        return rv;
 }
 
        struct driver_data *dd = pci_get_drvdata(pdev);
        int counter = 0;
 
+       set_bit(MTIP_DD_FLAG_REMOVE_PENDING_BIT, &dd->dd_flag);
+
        if (mtip_check_surprise_removal(pdev)) {
-               while (atomic_read(&dd->drv_cleanup_done) == false) {
+               while (!test_bit(MTIP_DD_FLAG_CLEANUP_BIT, &dd->dd_flag)) {
                        counter++;
                        msleep(20);
                        if (counter == 10) {
                return -EFAULT;
        }
 
-       atomic_set(&dd->resumeflag, true);
+       set_bit(MTIP_DD_FLAG_RESUME_BIT, &dd->dd_flag);
 
        /* Disable ports & interrupts then send standby immediate */
        rv = mtip_block_suspend(dd);
                dev_err(&pdev->dev, "Unable to resume\n");
 
 err:
-       atomic_set(&dd->resumeflag, false);
+       clear_bit(MTIP_DD_FLAG_RESUME_BIT, &dd->dd_flag);
 
        return rv;
 }